POJ 1741 Tree 樹上點對距離≤K計數 樹分治法

問題描述

給定一棵 n 個節點的樹（n ≤ 10,000），邊有權重。求節點對 (u, v) 使得 u 和 v 間距離 ≤ K 的數量。

樹分治核心思想

選擇當前子樹的重心作為根（刪除重心後最大子樹節點數最小），遞歸處理：
1. 統計經過重心的合法點對
2. 遞歸處理各子樹（排除重心）

經過重心的點對統計

• DFS 計算所有節點到重心的距離，存入數組 dist
• 排序 dist，雙指針法計算兩數之和 ≤ K 的數對
• 為避免重複計算同一子樹內的點對，對每個子樹單獨計算其內部合法點對並減去

關鍵代碼框架

int get_centroid(int u, int fa, int total_size) {
    size[u] = 1; max_sub[u] = 0;
    for (v : adj[u]) if (v != fa && !vis[v]) {
        get_centroid(v, u, total_size);
        size[u] += size[v];
        max_sub[u] = max(max_sub[u], size[v]);
    }
    max_sub[u] = max(max_sub[u], total_size - size[u]);
    if (max_sub[u] < max_sub[centroid]) centroid = u;
}

void get_dist(int u, int fa, int d, vector& dist) {
    dist.push_back(d);
    for (auto [v, w] : adj[u]) if (v != fa && !vis[v])
        get_dist(v, u, d + w, dist);
}

int count_pairs(vector& dist, int K) {
    sort(dist.begin(), dist.end());
    int ans = 0, l = 0, r = dist.size() - 1;
    while (l < r) {
        if (dist[l] + dist[r] <= K) ans += r - l, ++l;
        else --r;
    }
    return ans;
}

void solve(int u) {
    get_centroid(u, -1, subtree_size(u));
    int c = centroid; vis[c] = true;
    
    vector all_dist; get_dist(c, -1, 0, all_dist);
    ans += count_pairs(all_dist, K);
    
    for (auto [v, w] : adj[c]) if (!vis[v]) {
        vector sub_dist; get_dist(v, c, w, sub_dist);
        ans -= count_pairs(sub_dist, K);
        solve(v);
    }
}

複雜度分析

• 每層分治處理所有節點一次：O(n log n)
• 排序與雙指針：O(n log n)
• 總時間：O(n log² n)，空間：O(n)

關鍵優化

必須選擇重心分治，否則樹退化為鏈時遞歸深度 O(n) 導致 TLE 或棧溢出。