ผลรวมระยะทางในต้นไม้ใน C++

คำอธิบายปัญหา

ให้ต้นไม้ N โหนด (0..N-1) และ N-1 edges หา ans[i] = ผลรวมระยะทางจากโหนด i ไปทุกโหนดอื่น

ตัวอย่าง

N=6, edges=[[0,1],[0,2],[2,3],[2,4],[2,5]] → Output: [8,12,6,10,10,10]

แนวทางแก้ปัญหา: Tree DP (Two-pass DFS)

Pass 1 (Post-order): คำนวณจากใบขึ้นราก

• cnt[u] = จำนวนโหนดใน subtree u (รวม u)
• ans[u] = ผลรวมระยะทางจาก u ไปโหนดใน subtree u
• ans[u] = Σ (ans[v] + cnt[v]) สำหรับ child v

Pass 2 (Pre-order): คำนวณจากรากลงใบ (rerooting)

• เมื่อย้ายรากจาก u ไป child v: ans[v] = ans[u] - cnt[v] + (N - cnt[v]) (ระยะทางไปโหนดใน subtree v ใกล้ลง 1, ไปโหนดนอก subtree v ไกลขึ้น 1)

ตัวอย่างโค้ด C++

#include <bits/stdc++.h>
using namespace std;

class Solution {
    vector<int> ans, cnt;
    vector<vector<int>> graph;
    int N;
    
    void dfs1(int u, int p) {
        for(int v : graph[u]) if(v != p) {
            dfs1(v, u);
            cnt[u] += cnt[v];
            ans[u] += cnt[v] + ans[v];
        }
    }
    
    void dfs2(int u, int p) {
        for(int v : graph[u]) if(v != p) {
            ans[v] = ans[u] - cnt[v] + N - cnt[v];
            dfs2(v, u);
        }
    }
    
public:
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        this->N = N;
        ans.assign(N, 0);
        cnt.assign(N, 1);
        graph.assign(N, {});
        for(auto& e : edges) {
            graph[e[0]].push_back(e[1]);
            graph[e[1]].push_back(e[0]);
        }
        dfs1(0, -1);
        dfs2(0, -1);
        return ans;
    }
};

ความซับซ้อน

• เวลา: O(N)
• พื้นที่: O(N)