AtCoder abc220_f Distance Sums 2

题目链接

知识点：dp，dfs

题意

求树中每（一个节点到每一个节点的距离之和）。

思路

参考en_translator的题解。
$暴力O(n^2)，考虑节点之间的状态转移$
把一个节点当做根，对这个节点的一个相邻节点，找到任意一个节点，如果这个任意节点在这个相邻节点为根的子树中，从这个节点到这个相邻节点的状态转移时，这个任意节点的距离到目标距离会减$1$，在这个子树外则加$1$
$这个相邻节点的答案=这个节点的答案-这个子树的节点个数+这个子树外的节点个数$
$=这个节点的答案-这个子树的节点个数+总节点个数-这个子树的节点个数$
$=这个节点的答案+总节点个数-这个子树的节点个数\times 2$
实现思路见代码

代码

#include"cstdio"
#include"vector"
#include"bitset"
#include"algorithm"
#include"queue"

#define ll long long
#define vec vector<ll>
#define pb push_back
#define all(a) (a).begin(),(a).end()
const ll mod=998244353;
const ll N=2e6+10;
using namespace std;

vec graph[N];
bitset<N> vis;
ll ans[N];//见下文
ll cntsub[N];//i节点为根的子树节点个数

//计算节点1的答案和每个节点为根的子树的节点个数
void dfs(ll now,ll fa) {
  if(vis[now]) return;
  vis[now]=true;
  cntsub[now]=1;
  ans[now]=0;
  for(int i=0; i<graph[now].size(); i++) {
    ll nxt=graph[now][i];
    if(nxt==fa) continue;
    dfs(nxt,now);
    cntsub[now]+=cntsub[nxt];
    ans[now]+=ans[nxt]+cntsub[nxt];//nxt节点为根的子树的每个节点到目标的距离加1
    //ans[nxt]:原来的部分，cntsub[nxt]:加1的部分
  }
}

queue<ll> que;
ll n;

//计算所有节点的答案
void bfs() {
  vis.reset();
  while(!que.empty()) {
    ll now=que.front();
    que.pop();
    vis[now]=true;
    for(int i=0; i<graph[now].size(); i++) {
      ll nxt=graph[now][i];
      if(vis[nxt]) continue;
      que.push(nxt);
      ans[nxt]=ans[now]+n-cntsub[nxt]*2;
    }
  }
}

void init() {
  for(int i=1; i<=n; i++) graph[i].resize(0);
}

int main() {
  scanf("%lld",&n);
  init();
  for(int i=1; i<n; i++) {
    ll u,v;
    scanf("%lld%lld",&u,&v);
    graph[u].pb(v);
    graph[v].pb(u);
  }
  dfs(1,0);
  que.push(1);
  bfs();
  for(int i=1; i<=n; i++)
    printf("%lld\n",ans[i]);
  return 0;
}