(POJ) 1741.Tree [樹分治]

http://poj.org/problem?id=1741

題目大意：給一個有權重的樹，找出所有的pair(u,v)，使得在樹上u到v的最短路徑<=k。

這題主要是用到樹分治的技巧，在這之前要知道  樹重心  怎麼求。

樹分治的技巧在於：找到一個點，求出各個子樹的解之後，在合併起來。

那用樹重心的意義是說，這樣遞迴深度只有lgN層，整體複雜度就可以是O( N lgN ~~~)

那，這題，所有的解答可以分成三種case (1)在某個子樹T裡面 (2)在兩個不同的子樹 (3)一個子樹到樹重心

對於(1)，繼續遞迴下去求解即可

對於(2)、(3)，掃過這棵樹，把邊sort之後，認真維護一下

詳細可以看code喔~~~

練習題：Codeforces 715C. Digit Tree

#include <iostream>
#include <stdio.h>
#include <vector>
#include <utility>
#include <algorithm>
using namespace std;

typedef pair<int,int> pii;
const int MAX_N = 1e4 + 3;
const int INF = 1e9+7;

vector<pii> edg[MAX_N]; //to & weight
int tree[MAX_N];
pii node[MAX_N];  //sum & max
bool visit[MAX_N];
pii path[MAX_N];
pii p[MAX_N];
int cnt[MAX_N];

int dfs1(int id,int& tr) {
    tree[tr++]=id;
    visit[id]=true;
    int ret=1;
    int mx=0;
    for (vector<pii>::iterator i=edg[id].begin();i!=edg[id].end();i++) {
        pii tmp=*i;
        if (visit[tmp.first] == false) {
            int temp = dfs1(tmp.first,tr);
            ret += temp;
            mx = max(mx,temp);
        }
    }
    node[id]=make_pair(ret,mx);
    return ret;
}

int get_cen(int id) {
    int tr=0;
    dfs1(id,tr);
    int sz=tr;
    int ret_mx=sz;
    int ret=-1;
    for (int i=0;tr>i;i++) {
        int tmp=tree[i];
        int mx = max(node[tmp].second,sz - node[tmp].first);
        if (ret_mx > mx) {
            ret = tmp;
            ret_mx=mx;
        }
        visit[tmp]=false;
    }
    return ret;
}

void dfs2(int id,int p,int t,int &tr) {
    tree[tr++]=id;
    visit[id]=true;
    path[id]=make_pair(p,t);
    for (vector<pii>::iterator i=edg[id].begin();i!=edg[id].end();i++) {
        pii tmp=*i;
        if (visit[tmp.first] == false) {
            dfs2(tmp.first,p+tmp.second,t,tr);
        }
    }
}

int solve(int root,int k) {
    int ret=0;
    int cen = get_cen(root);
    visit[cen]=true;
    //case 1 --- all in subtree
    for (vector<pii>::iterator iter=edg[cen].begin();iter!=edg[cen].end();iter++) {
        pii tmp=*iter;
        if (visit[tmp.first]==false) ret += solve(tmp.first,k);
    }
    //case 2---cen & 3---two other subtree
    //case 2
    path[cen]=make_pair(0,0);
    int tree_id=0;
    int pid=0;
    for (vector<pii>::iterator iter=edg[cen].begin();iter!=edg[cen].end();iter++) {
        pii tmp=*iter;
        int tr=0;
        cnt[tree_id]=0;
        if (visit[tmp.first]==false) dfs2(tmp.first,tmp.second,tree_id,tr);
        for (int i=0;tr>i;i++) {
            int temp=tree[i];
            visit[temp]=false;
            if (path[temp].first <= k) {
                ret++;
                p[pid++] = (path[temp]);
                cnt[tree_id]++;
            } 
        }
        tree_id++;
    }
    //case 3
    sort(p,p+pid);
    int tmpans=0;
    int len=pid;
    int j=len-1;
    for (int i=0;len>i;i++) {
        while (j>=0 && p[j].first+p[i].first>k) {
            cnt[p[j].second]--;
            j--;
        }
        if (j<0) break;
        tmpans += (j+1-cnt[p[i].second]);
    }
    ret += tmpans/2;
    visit[cen]=false;
    return ret;
}

int main() {
//    freopen("input.txt","r",stdin);
    int n,k;
    while (scanf("%d %d",&n,&k) != EOF) {
        if (!n && !k) return 0;
        for (int x=0;n>=x;x++) {
            edg[x].clear();
            visit[x]=false;
        }
        for (int x=0;n-1>x;x++) {
            int i,j,k;
            scanf("%d %d %d",&i,&j,&k);
            edg[i].push_back(make_pair(j,k));
            edg[j].push_back(make_pair(i,k));
        }
        printf("%d\n",solve(1,k));
    }
}