POJ2114(树分治)

题意是给一棵最多1W个节点的树，求是否存在点对使得他们的路径是k。

类似于POJ的1741，修改下统计函数就可以了。注释掉的是一开始写的，很挫卡卡过，改掉以后效率好了一点。

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <iostream>
#include <cmath>
using namespace std;
#define maxn 111111
#define maxm 211111
#define INF 11111111
#define size Size

struct node {
    int from, to, next;
    int w;
}edge[maxm];
int n, m, cnt, root;
long long ans, k;
int head[maxn];
bool vis[maxn];
int size[maxn], num[maxn]; //以i为根的子树的节点 i的子树中最多的节点
long long d[maxn], Min;

void add_edge (int from, int to, int w) {
    edge[cnt].from = from, edge[cnt].to = to, edge[cnt].w = w, edge[cnt].next = head[from],
    head[from] = cnt++;
}

void dfs_size (int u, int fa) {
    size[u] = 1;
    num[u] = 0;
    for (int i = head[u]; i != -1; i = edge[i].next) {
        int v = edge[i].to;
        if (v == fa || vis[v])
            continue;
        dfs_size (v, u);
        size[u] += size[v];
        if (size[v] > num[u]) {
            num[u] = size[v];
        }
    }
}

void find_root (int u, int fa, int pre) { //找到重心
    long long cur = 0;
    cur = max (num[u], size[pre]-size[u]);
    if (cur < Min) {
        Min = cur;
        root = u;
    }
    for (int i = head[u]; i != -1; i = edge[i].next) {
        int v = edge[i].to;
        if (v == fa || vis[v])
            continue;
        find_root (v, u, pre);
    }
}

int tmp[maxn], tot;
void get_dis (int u, int fa, int dis) { //得到每个点到当前根的距离
    tmp[tot++] = dis;
    for (int i = head[u]; i != -1; i = edge[i].next) {
        int v = edge[i].to;
        if (vis[v] || v == fa)
            continue;
        d[v] = dis+edge[i].w;
        get_dis (v, u, d[v]);
    }
}

int solve (int u, int d) {
    long long ans = 0;
    tot = 0;
    get_dis (u, 0, d);
    sort (tmp, tmp+tot);
    int r = tot-1, l = 0;
    while (r>l) {
        for (; tmp[r]+tmp[l] > k && r > l; r--) {}
        /*if (tmp[l]+tmp[r] == k) {
            for (int i = r; i > l && tmp[l]+tmp[i] == k; i--)
                ans++;
        }*/
        if (tmp[l]+tmp[r] == k) {
            int cnt1 = 1, cnt2 = 1;
            while (r > l && tmp[r-1] == tmp[r]) {
                r--;
                if (r > l) cnt2++;
                else break;
            }
            while (r > l && tmp[l+1] == tmp[l]) {
                l++;
                if (r > l) cnt1++;
                else break;
            }
            ans += cnt1*cnt2;
        }
        l++;
    }
    return ans;
}

void dfs (int u) {
    Min = INF;
    dfs_size (u, 0);
    find_root (u, 0, u);
    vis[root] = 1;
    ans += solve (root, 0);
    for (int i = head[root]; i != -1; i = edge[i].next) {
        int v = edge[i].to;
        if (vis[v])
            continue;
        ans -= solve (v, edge[i].w);
        dfs (v);
    }
}

int main () {
    //freopen ("in", "r", stdin);
    while (scanf ("%d", &n) == 1 && n) { //cout << n << endl;
        cnt = 0;
        int u, v;
        long long w;
        memset (head, -1, sizeof head);
        for (u = 1; u <= n; u++) {
            for (; ;) {
                scanf ("%d", &v);
                if (!v)
                    break;
                scanf ("%lld", &w); //cout << u << " " << v << " " <<w << endl;
                add_edge (u, v, w);
                add_edge (v, u, w);
            }
        }
        for (; ;) {
            scanf ("%lld\n", &k);
            if (k == 0)
                break;
            memset (vis, 0, sizeof vis);
            ans = 0;
            dfs (1);
            if (ans) printf ("AYE\n");
            else printf ("NAY\n");
        }
        printf (".\n");
        //cout << k << endl;
    }
    //printf (".\n");
    return 0;
}
/*
4
2 0 3 0 0
4 0
0
0
1
2
0

1
0
1
2
3
0
*/