【算法练习】Luogu 3806 【模板】点分治1（点分治）

题意

给定一棵有n个点的树
询问树上距离为k的点对是否存在。

题解

点分治

代码

#include
using namespace std;
typedef double db;
typedef long long ll;
typedef unsigned long long ull;
const int nmax = 1e6+7;
const int INF = 0x3f3f3f3f;
const ll LINF = 0x3f3f3f3f3f3f3f3f;
const ull p = 67;
const ull MOD = 1610612741;
const int K = 1e7 + 7;
struct edge {
    int to, nxt, w;
}e[nmax << 1];
int head[nmax], tot;
int n, m;
void add_edge (int u, int v, int w) {
    e[tot].to = v;
    e[tot].nxt = head[u];
    e[tot].w = w;
    head[u] = tot++;
}
int rt ,nowmn = INF, id, subtreesize;
int sz[nmax], fa[nmax], dist[nmax], num[K];
bool visit[nmax];
void getroot(int u, int f) {
    sz[u] = 1;
    fa[u] = f;
    int mxchild = 0;
    for(int i = head[u]; i != -1; i = e[i].nxt ) {
        int v = e[i].to;
        if(v != f && !visit[v]) {
            getroot(v, u);
            sz[u] += sz[v];
            mxchild = max(mxchild, sz[v]);
        }
    }
    int tmp = max(mxchild, subtreesize - sz[u]);
    if(tmp < nowmn) {
        nowmn = tmp;
        rt = u;
    }
}
void getdis(int u, int f, int dis) {
    dist[id++] = dis;
    for(int i = head[u]; i != -1; i = e[i].nxt) {
        int v = e[i].to;
        if(v != f && !visit[v])
            getdis(v, u, dis + e[i].w);
    }
}
void getans(int u, int initdis, int tag) {
    id = 0;
    getdis(u, u, 0);
    for(int i = 0; i < id; ++i)
        for(int j = 0; j < id; ++j)
            if(tag == 1 && dist[i] + dist[j]<= K)
                num[dist[i] + dist[j]] += tag;
            else if (dist[i] + dist[j] + initdis <= K)
                num[dist[i] + dist[j] + initdis] += tag;
}

void solve(int u) {
    getans(u, 0, 1);
    visit[u] = true;
    for (int i = head[u]; i != -1; i = e[i].nxt) {
        int v = e[i].to;
        if (!visit[v]) {
            getans(v, e[i].w * 2, -1);
            nowmn = INF; rt = 0; subtreesize = sz[v];
            getroot(v, u);
            solve(rt);
        }
    }
}
int main(){
    memset(head, -1, sizeof head);
    scanf("%d %d", &n, &m);
    int u ,v, w;
    for(int i = 1; i <= n-1; ++i) {
        scanf("%d %d %d", &u, &v, &w);
        add_edge(u, v, w);
        add_edge(v, u, w);
    }
    subtreesize = n; nowmn = INF;
    getroot(1, -1);
    solve(1);
    int q;
    for(int i = 1; i <= m; ++i) {
        scanf("%d", &q);
        if(num[q])
            printf("AYE\n");
        else
            printf("NAY\n");
    }
    return 0;
}