POJ Tree：樹上距離不超過 K 的點對計數（點分治）

問題描述

給定 n 節點樹（n ≤ 10,000），邊有長度。求節點對 (u,v) 使得 dist(u,v) ≤ k 的數量。多測試用例，以 0 0 結束。

點分治核心思想

1. 找樹的重心：刪除後最大子樹最小。
2. 以重心為根，統計所有經過重心的合法點對（不同子樹間）。
3. 遞歸處理各子樹（同子樹內點對在遞歸中統計，避免重複）。

關鍵步驟

• 獲取重心：DFS 計算子樹大小，找最大分支最小的節點。
• 收集距離：從重心向各子樹 DFS，記錄到重心的距離 ≤ k 的節點距離值。
• 計數合法點對：對距離數組排序，雙指針法計算兩數之和 ≤ k 的對數（不同子樹）。同子樹內用相同邏輯但結果相減（排除）。

代碼框架

void getDist(int u, int fa, int dis) {
    if (dis > k) return;
    dist[tot++] = dis;
    for (auto e = head[u]; e; e = e->next)
        if (e->v != fa && !vis[e->v])
            getDist(e->v, u, dis + e->len);
}

void countPairs(int u) { // 不同子樹
    sort(dist, dist + tot);
    int l = 0, r = tot - 1;
    while (l < r) {
        if (dist[l] + dist[r] <= k) ans += r - l, ++l;
        else --r;
    }
}

void solve(int u) {
    int centroid = getCentroid(u);
    tot = 0; getDist(centroid, 0, 0); countPairs(centroid);
    vis[centroid] = 1;
    for (auto e = head[centroid]; e; e = e->next)
        if (!vis[e->v]) {
            tot = 0; getDist(e->v, centroid, e->len);
            sort(dist, dist + tot);
            int l = 0, r = tot - 1;
            while (l < r) { // 同子樹內減去
                if (dist[l] + dist[r] <= k) ans -= r - l, ++l;
                else --r;
            }
            solve(e->v);
        }
}

複雜度

每層分治 O(n log n)，深度 O(log n)，總體 O(n log² n)，n=10,000 輕鬆通過。