C++-এ ট্রি-এ দূরত্বের সমষ্টি

সমস্যা বিবরণ

N নোডের ট্রি দেওয়া থাকলে (এজ: edges[i] = [u,v])। প্রতিটি নোড i-এর জন্য অন্য সব নোড পর্যন্ত দূরত্বের সমষ্টি ans[i] বের করতে হবে।

উদাহরণ

N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]] → আউটপুট = [8,12,6,10,10,10]

সমাধান পদ্ধতি

• দুটি DFS পাস (Rerooting DP)।
• dfs1: বটম-আপ, cnt[node] = সাবট্রি সাইজ, ans[node] = সাবট্রি নোডগুলোর দূরত্ব সমষ্টি।
• dfs2: টপ-ডাউন, প্যারেন্ট থেকে চাইল্ডের ans আপডেট: ans[child] = ans[parent] - cnt[child] + (N - cnt[child])।
• সময় জটিলতা: O(N)।

C++ বাস্তবায়ন

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

void print_vector(vector<int> v) {
    cout << "[";
    for (int i = 0; i < v.size(); i++) cout << v[i] << ", ";
    cout << "]" << endl;
}

class Solution {
    vector<int> ans, cnt;
    vector<int> graph[10005];
    int N;
    
    void dfs1(int node, int parent) {
        for (int child : graph[node]) {
            if (child != parent) {
                dfs1(child, node);
                cnt[node] += cnt[child];
                ans[node] += cnt[child] + ans[child];
            }
        }
    }
    
    void dfs2(int node, int parent) {
        for (int child : graph[node]) {
            if (child != parent) {
                ans[child] = ans[node] - cnt[child] + N - cnt[child];
                dfs2(child, node);
            }
        }
    }
    
public:
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>> &edges) {
        this->N = N;
        ans = vector<int>(N);
        cnt = vector<int>(N, 1);
        
        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;
    }
};

int main() {
    Solution ob;
    vector<vector<int>> v = {{0,1},{0,2},{2,3},{2,4},{2,5}};
    print_vector(ob.sumOfDistancesInTree(6, v));
    return 0;
}

ইনপুট ও আউটপুট

ইনপুট: {{0,1},{0,2},{2,3},{2,4},{2,5}}

আউটপুট: [8, 12, 6, 10, 10, 10]