poj 2114 點分治

又做了一道點分治入門題.... 比較卡人的地方是。。 給你一個排好序的序列， 讓你求出裏面任意兩個數字加起來等於k的二元組的個數。。 O(n)掃一遍。。

#include <iostream>
#include <algorithm>
#include <cstring>
#include <cstdio>
#include <vector>
using namespace std;
int n,k;
const int maxn = 10010;
const int INF = 0x3f3f3f3f;
struct node{
    int to,next,w;
}edge[maxn*2];
int tot,head[maxn];
bool vis[maxn];
int size[maxn];
int dis[maxn];
int minn;
void init(){
    tot = 0; memset(head, -1, sizeof(head));
}
void add_edge(int u, int v, int w){
    edge[tot].to = v; edge[tot].w = w;
    edge[tot].next = head[u]; head[u] = tot++;
}
int getsize(int u, int pre){
    size[u] = 1; 
    for(int i=head[u]; ~i; i=edge[i].next){
        int v = edge[i].to;
        if(v == pre || vis[v] ) continue;
        size[u] += getsize(v, u);
    }
    return size[u];
}
void getroot(int u, int pre, int totnum, int &root){
    int maxx = totnum - size[u];
    for(int i=head[u]; ~i; i=edge[i].next){
        int v = edge[i].to;
        if(v == pre || vis[v]) continue;
        getroot(v, u, totnum, root);
        maxx = max(maxx, size[v]);
    }
    if(maxx < minn) {minn = maxx, root = u;}
}
int st,ed;
void getdis(int u, int pre, int len){
   dis[st++] = len;
   for(int i=head[u]; ~i; i=edge[i].next){
       int v = edge[i].to;
       if(v == pre || vis[v]) continue;
       getdis(v, u, len + edge[i].w);
   }
}
int getans(int a, int b){
   int e = b-1;
   int res = 0;
   sort(dis+a, dis+b);
   for(int i=a; i<b&&i<=e;){
        if(dis[i] + dis[e] > k) e--;
        else if(dis[i] + dis[e] < k) i++;
        else {
            if(dis[i] == dis[e]){
                res += (e - i + 1)*(e-i)/2;
                break;
            }
            int t1 = i; int t2 = e;
            while(dis[t1] == dis[i]) t1++;
            while(dis[t2] == dis[e]) t2--;
            res += (t1-i)*(e-t2);
            i = t1,e = t2;
        }
   }
   return res;
}
int solve(int u){
    int totnum = getsize(u, -1);
    int root;
    minn = INF;
    getroot(u, -1, totnum, root);
    vis[root] = true;
    int ret = 0;
    for(int i=head[root]; ~i; i=edge[i].next){
        int v = edge[i].to;
        if(vis[v]) continue;
        ret += solve(v); 
    }
    st = ed = 0;
    for(int i=head[root]; ~i; i=edge[i].next){
        int v = edge[i].to;
        if(vis[v]) continue;
        getdis(v, -1, edge[i].w); 
        ret -= getans(ed, st);
        ed = st;
    }
    ret += getans(0, ed);
    for(int i=0; i<ed; i++){
        if(dis[i] == k) ret ++;
        else if(dis[i] > k) break;
    }
    vis[root] = false;
    return ret;
}
int main(){
    int u,v,w;
    while(cin>>n&&n){
        init();
        for(int i=1; i<=n; i++){
            while(scanf("%d", &v) && v){
                scanf("%d", &w);
                add_edge(i, v, w);
                add_edge(v, i, w);
            }
        }
        while(scanf("%d", &k) == 1&&k){
            memset(vis, false, sizeof(vis));
            int res = solve(1);
            if(res) puts("AYE");
            else puts("NAY");
        }
        puts(".");
    }
    return 0;
}