洛谷【P3806】点分治

#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e4 + 5;
const int maxk = 1e7 + 5;
vector<pair<int, int> >edge[maxn];
///n点m询问,rt->当前重心,sum->当前树大小,cnt->计数器
int n, m, rt, sum, cnt;
///tmp->记录算出的距离,siz->记录子树大小,dis[i]->为rt与i之间的距离
///maxp->用于找重心,q->用于记录所有询问
int tmp[maxn], siz[maxn], dis[maxn], maxp[maxn], q[105];
///judge[i]->记录在之前子树中距离i是否存在,ans->记录第k个询问是否存在,vis->记录被删除的结点
bool judge[maxk], ans[105], vis[maxn];

inline int read(){
    int s = 0, w = 1; char ch = getchar();
    while(ch < '0' || ch > '9'){ if(ch == '-') w = -1; ch = getchar();}
    while(ch >= '0' && ch <= '9') s = s * 10 + ch - '0', ch = getchar();
    return s * w;
}

inline void print(int x){
    if(x < 0) x = ~x + 1, putchar('-');
    if(x > 9) print(x / 10);
    putchar(x % 10 + '0');
}

void dfs(int u, int f){///找重心
    siz[u] = 1, maxp[u] = 0;
    for(int i = 0; i < edge[u].size(); i++){
        int v = edge[u][i].first;
        if(v == f || vis[v]) continue;
        dfs(v, u);
        siz[u] += siz[v];
        if(siz[v] > maxp[u]) maxp[u] = siz[v];
    }
    maxp[u] = max(maxp[u], sum - siz[u]);
    if(maxp[u] < maxp[rt]) rt = u;
}

void dfs1(int u, int f){///统计当前子树的所有节点到根节点的距离
    tmp[cnt++] = dis[u];
    for(int i = 0; i < edge[u].size(); i++){
        int v = edge[u][i].first;
        if(v == f || vis[v]) continue;
        dis[v] = dis[u] + edge[u][i].second;
        dfs1(v, u);
    }
}

void solve(int u){///更新答案
    queue<int>Q;
    for(int i = 0; i < edge[u].size(); i++){
        int v = edge[u][i].first;
        if(vis[v]) continue;
        cnt = 0;
        dis[v] = edge[u][i].second;
        dfs1(v, u);
        for(int j = 0; j < cnt; j++)
            for(int k = 0; k < m; k++)
                if(q[k] >= tmp[j])
                    ans[k] |= judge[q[k] - tmp[j]];
        for(int j = 0; j < cnt; j++)
            if(tmp[j] < maxk) judge[tmp[j]] = true, Q.push(tmp[j]);
    }
    while(!Q.empty())
        judge[Q.front()] = false, Q.pop();
}

void divide(int u){///点分治
    vis[u] = judge[0] = true;
    solve(u);
    for(int i = 0; i < edge[u].size(); i++){
        int v = edge[u][i].first;
        if(vis[v]) continue;
        maxp[rt = 0] = sum = siz[v];
        dfs(v, 0);
        dfs(rt, 0);
        divide(rt);
    }
}

int main(){
    n = read(), m = read();
    for(int i = 1; i < n; i++){
        int u = read(), v = read(), w = read();
        edge[u].push_back(make_pair(v, w));
        edge[v].push_back(make_pair(u, w));
    }
    for(int i = 0; i < m; i++) q[i] = read();
    maxp[0] = sum = n;
    dfs(1, 0);
    dfs(rt, 0);
    divide(rt);
    for(int i = 0; i < m; i++){
        if(ans[i]) puts("AYE");
        else puts("NAY");
    }

    return 0;
}
/*
* ┌───┐   ┌───┬───┬───┬───┐ ┌───┬───┬───┬───┐ ┌───┬───┬───┬───┐ ┌───┬───┬───┐
* │Esc│   │ F1│ F2│ F3│ F4│ │ F5│ F6│ F7│ F8│ │ F9│F10│F11│F12│ │P/S│S L│P/B│  ┌┐    ┌┐    ┌┐
* └───┘   └───┴───┴───┴───┘ └───┴───┴───┴───┘ └───┴───┴───┴───┘ └───┴───┴───┘  └┘    └┘    └┘
* ┌───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───────┐ ┌───┬───┬───┐ ┌───┬───┬───┬───┐
* │~ `│! 1│@ 2│# 3│$ 4│% 5│^ 6│& 7│* 8│( 9│) 0│_ -│+ =│ BacSp │ │Ins│Hom│PUp│ │Num│ / │ * │ - │
* ├───┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─────┤ ├───┼───┼───┤ ├───┼───┼───┼───┤
* │ Tab │ Q │ W │ E │ R │ T │ Y │ U │ I │ O │ P │{ [│} ]│ | \ │ │Del│End│PDn│ │ 7 │ 8 │ 9 │   │
* ├─────┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴─────┤ └───┴───┴───┘ ├───┼───┼───┤ + │
* │ Caps │ A │ S │ D │ F │ G │ H │ J │ K │ L │: ;│" '│ Enter  │               │ 4 │ 5 │ 6 │   │
* ├──────┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴────────┤     ┌───┐     ├───┼───┼───┼───┤
* │ Shift  │ Z │ X │ C │ V │ B │ N │ M │< ,│> .│? /│  Shift   │     │ ↑ │     │ 1 │ 2 │ 3 │   │
* ├─────┬──┴─┬─┴──┬┴───┴───┴───┴───┴───┴──┬┴───┼───┴┬────┬────┤ ┌───┼───┼───┐ ├───┴───┼───┤ E││
* │ Ctrl│ Win│ Alt│         Space         │ Alt│ Win│Menu│Ctrl│ │ ← │ ↓ │ → │ │   0   │ . │←─┘│
* └─────┴────┴────┴───────────────────────┴────┴────┴────┴────┘ └───┴───┴───┘ └───────┴───┴───┘
*/