SPOJ1825 Free tour II (树分治)

/**
题意：N个节点的树，有M个节点是黑色，其余节点是白色，每条边有权值，要求找出一条路径，使得这条路径上黑色
节点数量在不超过K个的条件下长度最长

思路：设dp[x][k]:以x为根节点的且x为端点的所有路径上黑色节点数不超过k个时的最大长度，那么经过x的长度为l的
链的最大长度为dp[x][k] + dp[x][l - k](状态值需属于不同的子树)，但是这样复杂度太高，考虑树分治，只有logn层，
处理当前树时，答案路径要么在子树内，要么经过重心，子树的递归解决，设重心为x，经过重心的这样处理:
1. 处理出当前树所有节点vi，记录vi到x的路径长度dis和黑色节点数量col
2. 处理子树tree[i]的节点t时，有ans = max(ans, dis(t) + max{C[idx] | idx <= k - col})
C是树状数组，能够查询和更新区间[1, r]的最大值
3. 再用tree[i]的信息去更新树状数组
4. 处理tree[i + 1]

一层的复杂度是nlogn(排序，数组数组查询和更新)，所以总共是n * logn * logn
*/

#include<cstdio>
#include<cstring>
#include<cmath>
#include<vector>
#include<algorithm>
#include<iostream>
#include<set>
#include<queue>
#include<stack>
#include<functional>
typedef long long ll;
const int maxn = 3e5 + 10;
const int INF = 2e9 + 10;
using namespace std;

struct edge {
    int to, val;
    edge(int t = 0, int v = 0) : to(t), val(v) {}
};
struct cal {
    int x, dis, col;
    cal() {}
    cal(int x, int t, int f) : x(x), dis(t), col(f) {}
};
typedef pair<int, int> par;
int n, m, kase = 1, T, k, ans;
int siz[maxn], del[maxn], col[maxn];
vector<edge> G[maxn];
int C[maxn], tot_col[maxn];

void update(int x, int val, int tot) { for( ; x <= tot; x += x & -x) C[x] = max(C[x], val); }
int query_max(int x) { int ans = -INF; for( ; x; x -= x & -x) ans = max(ans, C[x]); return ans; }


int calculate_size(int x, int fa) {
    siz[x] = 1;
    for(int i = 0; i < G[x].size(); i++) {
        int v = G[x][i].to;
        if(v == fa || del[v]) continue;
        siz[x] += calculate_size(v, x);
    }
    return siz[x];
}

par find_zhongxin(int x, int fa, int tot) {
    int s = 1, m = 0;
    par res(INF, -1);
    for(int i = 0; i < G[x].size(); i++) {
        int v = G[x][i].to;
        if(v == fa || del[v]) continue;
        res = min(res, find_zhongxin(v, x, tot));
        s += siz[v]; m = max(m, siz[v]);
    }
    m = max(m, tot - s); res = min(res, par(m, x));
    return res;
}

void calculate_black_and_path(int x, int fa, int c, int d, vector<cal> &ds) {
    int co = c + col[x], dis = d;
    if(co > k) return ;
    ds.push_back(cal(x, dis, co));
    for(int i = 0; i < G[x].size(); i++) {
        int v = G[x][i].to, w = G[x][i].val;
        if(v == fa || del[v]) continue;
        calculate_black_and_path(v, x, co, dis + w, ds);
    }
}

int solve(int x) {
    calculate_size(x, 0);
    int zhongxin = find_zhongxin(x, 0, siz[x]).second, ans = 0;
    del[zhongxin] = 1;
    for(int i = 0; i < G[zhongxin].size(); i++) {
        int w = G[zhongxin][i].to;
        if(!del[w]) ans = max(ans, solve(w));
    }
    vector<cal> now;
    calculate_black_and_path(zhongxin, 0, 0, 0, now);
    int cnt = 0, need = k + col[zhongxin];
    for(int i = 0; i < now.size(); i++) {
        C[i] = C[i + 1] = 0;
        tot_col[cnt++] = now[i].col;
    }
    sort(tot_col, tot_col + cnt);
    cnt = unique(tot_col, tot_col + cnt) - tot_col;
    for(int j = 0; j < G[zhongxin].size(); j++) {
        int v = G[zhongxin][j].to, w = G[zhongxin][j].val;
        if(del[v]) continue;
        vector<cal> tds;
        calculate_black_and_path(v, zhongxin, col[zhongxin], w, tds);
        int tot = tds.size();
        for(int i = 0; i < tot; i++) {
            int ncol = tds[i].col, dis = tds[i].dis;
            int idx = upper_bound(tot_col, tot_col + cnt, need - ncol) - tot_col;
            int can_dis = query_max(idx);
            ans = max(ans, dis + can_dis);
        }
        for(int i = 0; i < tot; i++) {
            int ncol = tds[i].col, dis = tds[i].dis;
            int idx = upper_bound(tot_col, tot_col + cnt, ncol) - tot_col;
            update(idx, dis, cnt);
        }
    }
    del[zhongxin] = 0;
    return ans;
}

int main() {
    while(scanf("%d %d %d", &n, &k, &m) != EOF) {
        int from, to, val, vt;
        for(int i = 0; i < maxn; i++) G[i].clear();
        memset(del, 0, sizeof del);
        memset(col, 0, sizeof col);
        while(m--) { scanf("%d", &vt); col[vt] = 1; }
        for(int i = 0; i < n - 1; i++) {
            scanf("%d %d %d", &from, &to, &val);
            G[from].push_back(edge(to, val));
            G[to].push_back(edge(from, val));
        }
        printf("%d\n", solve(1));
    }
    return 0;
}