Hướng giải của Đếm cặp gần nhau

Nhớ rằng hướng dẫn giải này chỉ nên sử dụng khi bế tắc, và tuyệt đối không nên sao chép mã nguồn kèm theo. Hãy tôn trọng tác giả bài tập và người viết hướng dẫn giải.
Nộp mã nguồn lời giải chính thức trước khi giải bài tập đó có thể khiến bạn bị ban.

Tác giả: kienadmin

Lời giải

Ý tưởng

Sử dụng Centroid Decomposition để đếm số cặp có khoảng cách \(\le K\).

Tại mỗi centroid \(C\):

  1. Duyệt từng subtree con của \(C\), thu thập khoảng cách từ \(C\) đến các đỉnh trong subtree.
  2. Với mỗi khoảng cách \(d\) trong subtree hiện tại, đếm số đỉnh từ các subtree trước có khoảng cách \(\le K-d\).
  3. Thêm các khoảng cách của subtree hiện tại vào mảng đếm tổng.

Dùng Fenwick Tree hoặc prefix sum để truy vấn \(\le K-d\) nhanh.

Độ phức tạp
  • Thời gian: \(O(N \log^2 N)\).
  • Bộ nhớ: \(O(N)\).
Code mẫu
#include <bits/stdc++.h>
using namespace std;

const int MAXN = 20005;
vector<int> adj[MAXN];
int subsize[MAXN];
bool removed[MAXN];
int n, k;
long long ans = 0;
int bit[MAXN];

void add(int idx, int val) {
    // Cập nhật Fenwick Tree
    idx++;
    while (idx <= n + 1) {
        bit[idx] += val;
        idx += idx & -idx;
    }
}

int sum(int idx) {
    // Truy vấn tổng từ 0 đến idx
    idx++;
    int res = 0;
    while (idx > 0) {
        res += bit[idx];
        idx -= idx & -idx;
    }
    return res;
}

void get_dists(int u, int p, int d, vector<int>& dists) {
    // Thu thập khoảng cách từ centroid đến các đỉnh trong subtree
    dists.push_back(d);
    for (int v : adj[u]) {
        if (v != p && !removed[v]) {
            get_dists(v, u, d + 1, dists);
        }
    }
}

void count_at_centroid(int u) {
    // Đếm số cặp đi qua centroid u
    for (int v : adj[u]) {
        if (removed[v]) continue;
        vector<int> dists;
        get_dists(v, u, 1, dists);
        for (int d : dists) {
            if (d <= k) ans += sum(k - d);
        }
        for (int d : dists) {
            if (d <= k) add(d, 1);
        }
    }
    // Reset BIT: chỉ clear những ô đã dùng
    for (int v : adj[u]) {
        if (removed[v]) continue;
        vector<int> dists;
        get_dists(v, u, 1, dists);
        for (int d : dists) if (d <= k) add(d, -1);
    }
}

void dfs_size(int u, int p) {
    subsize[u] = 1;
    for (int v : adj[u])
        if (v != p && !removed[v]) {
            dfs_size(v, u);
            subsize[u] += subsize[v];
        }
}

int find_centroid(int u, int p, int total) {
    for (int v : adj[u])
        if (v != p && !removed[v] && subsize[v] > total / 2)
            return find_centroid(v, u, total);
    return u;
}

void decompose(int u) {
    dfs_size(u, -1);
    int c = find_centroid(u, -1, subsize[u]);
    // Thêm đỉnh centroid vào BIT (khoảng cách 0)
    add(0, 1);
    count_at_centroid(c);
    add(0, -1);
    removed[c] = true;
    for (int v : adj[c])
        if (!removed[v]) decompose(v);
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cin >> n >> k;
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    decompose(1);
    cout << ans << "\n";
    return 0;
}

Nhận xét

Không có ý kiến tại thời điểm này.