【模板】点分治

点分治

POJ-1741

#include <cstdio>
#include <algorithm>
using namespace std;
const int N = 10010;
const int INF = 0x3f3f3f3f;

int n, k, ans;

int head[N], nxt[2 * N], to[2 * N], w[2 * N], tot; //链式前向星
void add_edge(int x, int y, int z) {
    nxt[++tot] = head[x];
    to[tot] = y;
    w[tot] = z;
    head[x] = tot;
}

bool vis[N];
int Root, Tsiz, siz[N], mx[N];
int arr[N], cnt;

void init() {
    tot = ans = 0;
    for (int i = 0; i < N; i++) vis[i] = 0, head[i] = -1;
    Tsiz = n; mx[0] = INF; //最大子树节点数mx[0]初始化为INF
}

void GetRoot(int u, int fa) { //寻找树的重心
    siz[u] = 1; mx[u] = 0;
    for (int i = head[u]; ~i; i = nxt[i]) {
        if (to[i] != fa && !vis[to[i]]) {
            GetRoot(to[i], u);
            siz[u] += siz[to[i]];
            mx[u] = max(mx[u], siz[to[i]]);
        }
        mx[u] = max(mx[u], Tsiz - siz[u]);
        if (mx[Root] > mx[u]) Root = u;
    }
}

void GetDis(int u, int D, int fa) { //统计路径长度
    arr[++cnt] = D;
    for (int i = head[u]; ~i; i = nxt[i]) {
        if (to[i] != fa && !vis[to[i]])
            GetDis(to[i], D + w[i], u);
    }
}

int calc(int u, int D) {
    cnt = 0;
    GetDis(u, D, 0);
    int l = 1, r = cnt, sum = 0;
    sort(arr + 1, arr + 1 + cnt);
    for (;;++l) {
        while (r && arr[l] + arr[r] > k) --r;
        if (r < l) break;
        sum += r - l + 1;
    }
    return sum;
}

void DAC(int u) { //分治，u为当前树的重心
    ans += calc(u, 0); //加上当前点的答案
    vis[u] = 1;
    for (int i = head[u]; ~i; i = nxt[i]) {
        if (!vis[to[i]]) {
            ans -= calc(to[i], w[i]); //容斥原理
            Root = 0, Tsiz = siz[to[i]];
            GetRoot(to[i], 0); //得到子树的重心
            DAC(Root); //递归求解子树
        }
    }
} //此法时间复杂度不会退化

int main() {
    while (~scanf("%d %d", &n, &k) && n && k) {
        init();
        for (int i = 1; i < n; i++) {
            int x, y, z; scanf("%d %d %d", &x, &y, &z);
            add_edge(x, y, z); add_edge(y, x, z);
        }
        GetRoot(1, 0);
        DAC(Root);
        printf("%d\n", ans - n); //去除(u, u)点对
    }
    return 0;
}

luogu_P3806

#include <cstdio>
#include <algorithm>
#include <cstring>
using namespace std;
const int N = 10010; 
const int INF = 0x3f3f3f3f; 

int n, m;
int query[N];
bool vis[N], ok[N];
int e_cnt, head[N], mx[N], sz[N], Root, cnt, a[N], b[N], d[N];
/* 记当前分治的根为Root
a[]记录从Root能到的点
d[]记录a_i到Root的距离
b[]记录a_i属于Root的哪一个子树（当b[a[i]] == b[a[j]]时，说明a_i与a_j属于Root的同一颗子树
*/
struct Edge
{
    int to, nxt, w;
}edge[N << 1];

void add_edge(int x, int y, int z) {
    edge[++e_cnt].nxt = head[x];
    edge[e_cnt].to = y;
    edge[e_cnt].w = z;
    head[x] = e_cnt;
}

void getRoot(int u, int fa, int total) {
    sz[u] = 1, mx[u] = 0;
    for (int i = head[u]; ~i; i = edge[i].nxt) {
        int v = edge[i].to;
        if (v == fa || vis[v]) continue;
        getRoot(v, u, total);
        sz[u] += sz[v];
        mx[u] = max(mx[u], sz[v]);
    }
    mx[u] = max(mx[u], total - sz[u]);
    if (mx[u] < mx[Root]) Root = u;
}

void getDis(int u, int fa, int dis, int from) {
    a[++cnt] = u;
    d[u] = dis;
    b[u] = from;
    for (int i = head[u]; ~i; i = edge[i].nxt) {
        int v = edge[i].to;
        if (v == fa || vis[v]) continue;
        getDis(v, u, dis + edge[i].w, from);
    }
}

bool cmp(int x, int y) { return d[x] < d[y]; }

void calc(int u) {
    cnt = 0;
    a[++cnt] = u;
    d[u] = 0;
    b[u] = u;
    for (int i = head[u]; ~i; i = edge[i].nxt) {
        int v = edge[i].to;
        if (vis[v]) continue;
        getDis(v, u, edge[i].w, v);
    }
    sort(a + 1, a + 1 + cnt, cmp);
    for (int i = 1; i <= m; i++) {
        int l = 1, r = cnt;
        if (ok[i]) continue;
        while (l < r) {
            if (d[a[l]] + d[a[r]] > query[i]) r--;
            else if (d[a[l]] + d[a[r]] < query[i]) l++;
            else if (b[a[l]] == b[a[r]]) {
                if (d[a[r]] == d[a[r - 1]]) r--;
                else l++;
            }
            else { ok[i] = 1; break; }
        }
    }
}

void DAC(int u) {
    vis[u] = 1;
    calc(u);
    for (int i = head[u]; ~i; i = edge[i].nxt) {
        int v = edge[i].to;
        if (vis[v]) continue;
        Root = 0;
        getRoot(v, 0, sz[v]);
        DAC(Root);
    }
}

int main() {
    memset(head, -1, sizeof(head));
    scanf("%d %d", &n, &m);
    for (int i = 1; i < n; i++) {
        int x, y, z; scanf("%d %d %d", &x, &y, &z);
        add_edge(x, y, z); add_edge(y, x, z);
    }
    for (int i = 1; i <= m; i++) {
        scanf("%d", &query[i]);
        if (!query[i]) ok[i] = 1;
    }
    mx[0] = INF;
    getRoot(1, 0, n);
    DAC(Root);
    for (int i = 1; i <= m; i++) {
        if (ok[i]) printf("AYE\n");
        else printf("NAY\n");
    }
}