2023/7/16-leetcode每日一题-树形DP
https://leetcode.cn/problems/sum-of-distances-in-tree/
树形DP
class Solution {
int[] answer;
int[] sz;
int[] dp;
ArrayList<List<Integer>> graph;
public void dfs(int u,int f){
dp[u]=0;
sz[u]=1;
for(int v:graph.get(u)){
if(v==f){
continue;
}
dfs(v,u);
dp[u] += dp[v] + sz[v];
sz[u] += sz[v];
}
}
public void dfs2(int u, int f){
answer[u]=dp[u];
for(int v:graph.get(u)){
if(v==f){
continue;
}
else{
int du=dp[u];
int su=sz[u];
int dv=dp[v];
int sv=sz[v];
dp[u]-=dp[v]+sz[v];
sz[u]-=sz[v];
dp[v]+=dp[u]+sz[u];
sz[v]+=sz[u];
//不断深搜直到最深,没有子树为止,然后保存dp值,就是结果answer
dfs2(v,u);
//这是为下一个循环服务的,应该恢复
dp[u]=du;
dp[v]=dv;
sz[u]=su;
sz[v]=sv;
}
}
}
public int[] sumOfDistancesInTree(int n, int[][] edges) {
answer = new int[n];
sz = new int[n];
dp = new int[n];
graph = new ArrayList<List<Integer>>();
if(n==1){
answer[0] = 0;
return answer;
}
if(n==2){
answer[0] = 1;
answer[1] = 1;
return answer;
}
for(int i=0;i<n;i++){
graph.add(new ArrayList<Integer>());
}
for(int[] edge: edges){
int u=edge[0];
int v=edge[1];
graph.get(u).add(v);
graph.get(v).add(u);
}
dfs(0,-1);
dfs2(0,-1);
return answer;
}
}