SSL2541 2016提高组模拟试题 敌敌（树形dp）

2016提高组模拟试题（10月28日） 敌敌

Description

月考刚考了年级倒数的敌敌来到电脑室，朝着正在写树套树套树套树的beginend大吼：我就不信我有辣么辣鸡！我可是去年普及组AC了第一题的人啊！
beginend不耐烦地扔了一道题给敌敌，然后说：你只要把这道题写出来你就不是辣鸡啦。
敌敌接过题一看，题目是这样的：
有n个点和n-1条无向边，满足任意两点都可以互相到达，求有多少条经过且仅经过三个点的简单路径。
（简单路径即为每个点最多被访问一次的路径）
但是敌敌发现自己并不会写这题，但又不愿承认自己是辣鸡这个事实，于是他就拿着题目来找刚刚成为保送狗的你。为了敌敌的智商，你只好选择了帮他一把。

Input

第一行一个正整数n表示城市数量。
接下来n-1行每行两个正整数u,v表示u和v之间有一条道路。

Output

输出一行，表示经过且仅经过三个点的简单路径的数量。
Sample Input

3
1 2
2 3
Sample Output

2
Hint

【数据范围与约定】
对于 30%的数据 n<=20
对于 60%的数据 n<=200
对于 100%的数据 n<=200000

分析：由于只有3个点，考虑每个点对答案的贡献只与儿子和孙子有关，那显然设f[i][0]为i的儿子个数，f[i][1]为i的孙子个数，i的贡献就为f[i][0]*(f[i][0]-1)+f[i][1]*2。

代码

#include <cstdio>
#define maxn 500000
using namespace std;

struct arr
{
    int to,nxt;
}a[maxn];
int n,l,ls[maxn];
long long ans,f[maxn][2];

void add(int x,int y)
{
    a[++l].to=y;
    a[l].nxt=ls[x];
    ls[x]=l;
}

void dfs(int x,int fa)
{
    for (int i=ls[x];i;i=a[i].nxt)
        if (a[i].to!=fa)
        {
            f[x][0]++;
            dfs(a[i].to,x);
            f[x][1]+=f[a[i].to][0];
        }
    ans+=f[x][0]*(f[x][0]-1)/2+f[x][1];
}

int main()
{
    scanf("%d",&n);
    for (int i=1;i<n;i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);
        add(y,x);
    }
    dfs(1,0);
    printf("%lld",ans*2);
}