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

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

输入:

       3

      1  2

      1  3

输出:

　　0 2 2

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

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const LL mod=1e9+7;
const int N=1e5+7;
vector <int> g[N];
LL sum[N][3], num[N][3];
bool vis[N];
int n; LL ans[3];
LL my_add(int num1, int sum1, int num2, int sum2) {
    return sum1*num2%mod+sum2*num1%mod+num1*num2%mod;
}
void dfs(int rt) {
    vis[rt]=1;
    for (int i=0;i<g[rt].size();i++) {
        int child=g[rt][i];
        if (!vis[child]) {
            dfs(child);
            for (int i=0;i<3;i++)
                for (int j=0;j<3;j++)  {
                    int k=(i+j+1)%3;
                    ans[k]=(ans[k]+my_add(num[rt][i],sum[rt][i],num[child][j],sum[child][j]))%mod;
                }
            for (int i=0;i<3;i++) {
                num[rt][(i+1)%3]+=num[child][i];
                sum[rt][(i+1)%3]=(sum[rt][(i+1)%3]+(sum[child][i]+num[child][i])%mod)%mod;
            }
        }
    }
    num[rt][0]+=1; // importment!!!(别忘了自己)
}
int main()
{
    scanf("%d", &n);
    for (int i=1;i<n;i++) {
        int u,v; scanf("%d %d",&u,&v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1);
    for (int i=0;i<3;i++)
        for (int j=1;j<=n;j++)
        ans[i]=(ans[i]+sum[j][i])%mod;
    printf("%lld %lld %lld
",ans[0], ans[1], ans[2]);
    return 0;
}