HDU 5416 CRB and Tree (樹形dp)

題目：http://acm.hdu.edu.cn/showproblem.php?pid=5416

題意：給出一棵樹和每條邊的權值，f(u,v) 表示結點u到結點v之間路徑上所有邊異或運算的結果，q次詢問，每次問f(u,v)=s的路徑有多少條，u可以等於v。

代碼：

#include <stdio.h>
#include <string.h>
#include <vector>
using namespace std;
typedef long long ll;
int n, cnt[200005], vis[200005];
struct Edge
{
    int to, cost;
}e;
vector<Edge> G[100005];

void dfs(int v, int cur)
{//f(1,v)=cur
    vis[v] = 1;
    for(int i = 0; i < G[v].size(); i++)
    {
        e = G[v][i];
        if(!vis[e.to])
        {
            cnt[cur ^ e.cost]++;
            dfs(e.to, cur ^ e.cost);
        }
    }
}

int main()
{
    int t, a, b, c, q, s;
    scanf("%d", &t);
    while(t--)
    {
        scanf("%d", &n);
        for(int i = 0; i <= n; i++)
            G[i].clear();
        for(int i = 0; i < n - 1; i++)
        {
            scanf("%d%d%d", &a, &b, &c);
            e.to = b; e.cost = c;
            G[a].push_back(e);
            e.to = a;
            G[b].push_back(e);
        }
        memset(cnt, 0, sizeof(cnt)); //cnt[i]記錄f(1,u)=i的結點u的個數
        memset(vis, 0, sizeof(vis));
        dfs(1, 0); //樹形dp求cnt[]
        scanf("%d", &q);
        while(q--)
        {
            scanf("%d", &s);
            ll ans = 0;
            ans += cnt[s]; //f(1,u)=s的結點u個數（u不等於1）
            memset (vis, 0, sizeof(vis));
            if(s == 0)
            {
                ans += n; //每個結點u都有f(u,u)=0
                for(int i = 0; i <= ((1<<17)-1); i++)
                {
                    if(cnt[i]) //f(1,u)=f(1,v)則f(u,v)=0
                        ans += (ll)cnt[i] * (cnt[i] - 1) / 2;
                }
            }
            else
            {
                for(int i = 0; i <= ((1<<17)-1); i++)
                {//f(1,u)=i,f(1,v)=j則f(u,v)=i^j=s,j=s^i
                    if(!vis[i])
                    {
                        ans += (ll)cnt[i] * cnt[s ^ i];
                        vis[i] = vis[s ^ i] = 1;
                    }
                }
            }
            printf("%I64d\n", ans);
        }
    }
    return 0;
}