2020-字节跳动笔试(树距离之和[距离按%3值不同,分为三类])

题意描述, 有一颗树, 你要把树上任意不同两点的距离按%3的值不同分为3类...输出所有距离%3==k之和

输入:

1 2

1 3

输出:

　　0 2 2

分析: 和https://www.cnblogs.com/xidian-mao/p/11297683.html这道题思路类似, 只不过要分类去求(自己编程效率太慢了~~~~, 所有这是你不理我的原因嘛)

 1 #include <bits/stdc++.h>
 2 using namespace std;
 3 typedef long long LL;
 4 const LL mod=1e9+7;
 5 const int N=1e5+7;
 6 vector <int> g[N];
 7 LL sum[N][3], num[N][3];
 8 bool vis[N];
 9 int n; LL ans[3];
10 LL my_add(int num1, int sum1, int num2, int sum2) {
11     return sum1*num2%mod+sum2*num1%mod+num1*num2%mod;
12 }
13 void dfs(int rt) {
14     vis[rt]=1;
15     for (int i=0;i<g[rt].size();i++) {
16         int child=g[rt][i];
17         if (!vis[child]) {
18             dfs(child);
19             for (int i=0;i<3;i++)
20                 for (int j=0;j<3;j++)  {
21                     int k=(i+j+1)%3;
22                     ans[k]=(ans[k]+my_add(num[rt][i],sum[rt][i],num[child][j],sum[child][j]))%mod;
23                 }
24             for (int i=0;i<3;i++) {
25                 num[rt][(i+1)%3]+=num[child][i];
26                 sum[rt][(i+1)%3]=(sum[rt][(i+1)%3]+(sum[child][i]+num[child][i])%mod)%mod;
27             }
28         }
29     }
30     num[rt][0]+=1; // importment!!!(别忘了自己)
31 }
32 int main()
33 {
34     scanf("%d", &n);
35     for (int i=1;i<n;i++) {
36         int u,v; scanf("%d %d",&u,&v);
37         g[u].push_back(v);
38         g[v].push_back(u);
39     }
40     dfs(1);
41     for (int i=0;i<3;i++)
42         for (int j=1;j<=n;j++)
43         ans[i]=(ans[i]+sum[j][i])%mod;
44     printf("%lld %lld %lld
",ans[0], ans[1], ans[2]);
45     return 0;
46 }