poj2112 Boatherds 点分治

和上一道一样，只是把<=K改成=K就行了，不过算起来难度增加了一点，滑窗的时候容易滑错。

这次我把端点在根的也弄到第三种情况一起搞，也就是分两种情况，经过根和完全在子树里。这样写起来也之前写的简洁不少。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<vector>
#define REP(i,a,b) for(int i=a;i<=b;i++)
#define MS0(a) memset(a,0,sizeof(a))

using namespace  std;

typedef long long ll;
const int maxn=10100;
const int INF=1e9+10;

int n,k;
int u,v,w;
struct Edge
{
    int v,w;
};
vector<Edge> G[maxn];
vector<int> d;
bool vis[maxn];
int rt,balance;

void dfs_d(int u,int f,int dep)
{
    d.push_back(dep);
    for(int i=0;i<G[u].size();i++){
        int v=G[u][i].v,w=G[u][i].w;
        if(v==f||vis[v]) continue;
        dfs_d(v,u,dep+w);
    }
}

int calc(int u,int f,int dep)
{
    d.clear();
    dfs_d(u,f,dep);
    sort(d.begin(),d.end());
    int px=0,pl=-1;
    int res=0;
    for(int l=0,r=(int)d.size()-1;l<r;l++){
        if(d[l]==pl) res+=px;
        else{
            while(r>l&&d[l]+d[r]>k) r--;
            px=0;
            while(r>l&&d[l]+d[r]==k) r--,px++;
            if(r<(int)d.size()-1) r++;
            res+=px;
            pl=d[l];
            if(d[l]==d[r]&&d[l]+d[r]==k) res+=px*(px-1)/2;
        }
    }
    return res;
}

int get_rt(int u,int f,int sz)
{
    int cnt=1,balance1=0;
    for(int i=0;i<G[u].size();i++){
        int v=G[u][i].v;
        if(v==f||vis[v]) continue;
        int tmp=get_rt(v,u,sz);
        cnt+=tmp;
        balance1=max(balance1,tmp);
    }
    balance1=max(balance1,sz-cnt);
    if(balance1<balance){
        balance=balance1;
        rt=u;
    }
    return cnt;
}

int solve(int u)
{
    rt=u,balance=INF;
    int sz=get_rt(u,0,n);
    rt=u,balance=INF;
    get_rt(u,0,sz);
    u=rt;
    //cout<<"u="<<u<<endl;
    vis[u]=1;
    int res=0;
    res+=calc(u,0,0);
    //cout<<"res1="<<res<<endl;
    for(int i=0;i<G[u].size();i++){
        int v=G[u][i].v,w=G[u][i].w;
        if(vis[v]) continue;
        res-=calc(v,u,w);
    }
    //cout<<"res2="<<res<<endl;
    for(int i=0;i<G[u].size();i++){
        int v=G[u][i].v;
        if(vis[v]) continue;
        res+=solve(v);
    }
    return res;
}

void test()
{
    int n,a;
    cin>>n;
    d.clear();
    REP(i,1,n) cin>>a,d.push_back(a);
    sort(d.begin(),d.end());
    while(cin>>k){
        int res=0;
        int pl=-1,px=0;
        for(int l=0,r=(int)d.size()-1;l<r;l++){
            if(d[l]==pl) res+=px;
            else{
                while(r>l&&d[l]+d[r]>k) r--;
                px=0;
                while(r>l&&d[l]+d[r]==k) r--,px++;
                if(r<(int)d.size()-1) r++;
                res+=px;
                pl=d[l];
                if(d[l]==d[r]&&d[l]+d[r]==k) res+=px*(px-1)/2;
            }
        }
        cout<<res<<endl;
    }
}

int main()
{
    //test();
    freopen("in.txt","r",stdin);
    while(cin>>n,n){
        REP(i,1,n) G[i].clear();
        REP(u,1,n){
            while(scanf("%d",&v)&&v){
                scanf("%d",&w);
                G[u].push_back({v,w});
                G[v].push_back({u,w});
            }
        }
        while(scanf("%d",&k)&&k){
            MS0(vis);
            //printf("%d\n",solve(1));
            puts(solve(1)?"AYE":"NAY");
        }
        puts(".");
    }
    return 0;
}