洛谷P3806 【模板】点分治1 #点分治#

题目背景

感谢 hzwer 的点分治互测。

题目描述

给定一棵有 nn 个点的树。

询问树上距离为 kk 的点对是否存在。

输入格式

第一行两个数 n,mn,m。

接下来 n-1n−1 条边 a,b,ca,b,c 描述 aa 到 bb 有一条长度为 cc 的路径。

接下来 mm 行每行询问一个 KK。

输出格式

对于每个 KK 每行输出一个答案，存在输出 AYE，否则输出 NAY。

输入输出样例

输入 #1复制

2 1
1 2 2
2

输出 #1复制

AYE

说明/提示

对于 30\%30% 的数据，n\leq 100n≤100。

对于 60\%60% 的数据，n\leq 1000n≤1000，m\leq 50m≤50 。

对于 100\%100% 的数据，n\leq 10^4n≤104，m\leq 100m≤100，c\leq 10^4c≤104，K\leq 10^7K≤107。

题解

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<int, int> p;
const double pi = acos(-1.0);
const int inf = 0x3f3f3f3f;
const int maxn = 1e4 + 10;
const int maxm = 1e2 + 10;
const int maxk = 2e7 + 10;
bool vis[maxn], judge[maxk];
int tot, head[maxn];
int rt, sum, root[maxn], siz[maxn], maxp[maxn];
int cnt, tmp[maxn], dis[maxn];
int que[maxm], ans[maxm];
struct edge { int to, dis, nxt; } e[maxn << 1];

template<typename T = int>
inline const T read()
{
    T x = 0, f = 1; char ch = getchar();
    while (ch < '0' || ch > '9') { if (ch == '-') f = -1; ch = getchar(); }
    while (ch >= '0' && ch <= '9') { x = (x << 3) + (x << 1) + ch - '0'; ch = getchar(); }
    return x * f;
}

template<typename T>
inline void write(T x)
{
    if (x < 0) { putchar('-'); x = -x; }
    if (x > 9) write(x / 10);
    putchar(x % 10 + '0');
}

void addEdge(int u, int v, int w)
{
    e[tot].to = v;
    e[tot].dis = w;
    e[tot].nxt = head[u];
    head[u] = tot++;
}

void getRoot(int u, int f)
{
    siz[u] = 1;
    maxp[u] = 0;
    for (int i = head[u]; ~i; i = e[i].nxt)
    {
        int v = e[i].to;
        if (vis[v] || v == f) continue;
        getRoot(v, u);
        siz[u] += siz[v];
        maxp[u] = max(maxp[u], siz[v]);
    }
    maxp[u] = max(maxp[u], sum - siz[u]);
    if (maxp[u] < maxp[rt]) rt = u;
}

void getDis(int u, int f)
{
    tmp[cnt++] = dis[u];
    for (int i = head[u]; ~i; i = e[i].nxt)
    {
        int v = e[i].to, w = e[i].dis;
        if (vis[v] || v == f) continue;
        dis[v] = dis[u] + w;
        getDis(v, u);
    }
}

void solve(int u, int m)
{
    static queue<int> q;
    for (int i = head[u]; ~i; i = e[i].nxt)
    {
        int v = e[i].to;
        if (vis[v]) continue;
        cnt = 0;
        dis[v] = e[i].dis;
        getDis(v, u);
        for (int j = 0; j < cnt; j++)
            for (int k = 0; k < m; k++)
                if (que[k] >= tmp[j])
                    ans[k] |= judge[que[k] - tmp[j]];
        for (int j = 0; j < cnt; j++)
        {
            q.push(tmp[j]);
            judge[tmp[j]] = true;
        }
    }
    while (!q.empty())
    {
        judge[q.front()] = false;
        q.pop();
    }
}

void divide(int u, int m)
{
    vis[u] = judge[0] = true;
    solve(u, m);
    for (int i = head[u]; ~i; i = e[i].nxt)
    {
        int v = e[i].to;
        if (vis[v]) continue;
        maxp[rt = 0] = sum = siz[v];
        getRoot(v, 0);
        getRoot(rt, 0);
        divide(rt, m);
    }
}

int main()
{
#ifdef ONLINE_JUDGE
#else
    freopen("input.txt", "r", stdin);
#endif
    memset(head, -1, sizeof(head));
    int n = read(), m = read();
    for (int i = 0; i < n - 1; i++)
    {
        int u = read(), v = read(), w = read();
        addEdge(u, v, w);
        addEdge(v, u, w);
    }
    for (int i = 0; i < m; i++) que[i] = read();
    maxp[0] = sum = n;
    getRoot(1, 0);
    getRoot(rt, 0);
    divide(rt, m);
    for (int i = 0; i < m; i++) printf("%s\n", ans[i] ? "AYE" : "NAY");
    return 0;
}