洛谷 P3806 【模板】点分治1

题目背景

感谢 hzwer 的点分治互测。

题目描述

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

输入格式

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

第 22 到第 nn 行，每行三个整数 u, v, wu,v,w，代表树上存在一条连接 uu 和 vv 边权为 ww 的路径。

接下来 mm 行，每行一个整数 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% 的数据，保证 1 \leq n\leq 10^41≤n≤104，1 \leq m\leq 1001≤m≤100，1 \leq k \leq 10^71≤k≤107，1 \leq u, v \leq n1≤u,v≤n，1 \leq w \leq 10^41≤w≤104。

提示

• 本题不卡常。
• 如果您 #7 一直 RE/TLE，不妨看看 这个帖子。
  请注意，第一篇题解也存在帖子中说明的问题，但是鉴于此问题与点分治算法本身无关，并且题解质量非常高，因此暂不撤下题解，仅供参考。

模板来自AgOH

视频讲解

#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e4+5;
const int maxk = 2e7+5;
struct E {
    int to, w, next;
} edge[maxn << 1];
int tot, head[maxn];
inline void addedge(int u, int v, int w) {
    edge[tot] = (E) {
        v, w, head[u]
    };
    head[u] = tot++;
}
//? rt记录重心，all记录当前树大小，cnt是计数器
int n, m, rt, rt2, all, cnt;
//? tmp记录算出的距离，siz记录子树大小，dis[i]为rt与i之间的距离
//? maxson用于找重心，q用于记录所有询问
int tmp[maxn], siz[maxn], dis[maxn], maxson[maxn], q[105];
//? judge[i]记录在之前子树中距离i是否存在，ans记录第k个询问是否存在，vis记录被删除的结点
bool judge[maxk], ans[105], vis[maxn];

//TODO 找重心 最多两个: rt rt2
void getrt(int u, int fa) {
	siz[u] = 1;
	maxson[u] = 0;  //maxson初始化为最小值
	//遍历所有儿子，用maxson保存最大儿子的大小
	for(int i = head[u]; ~i; i = edge[i].next) {
		int v = edge[i].to;
		if(v == fa || vis[v]) continue; //被删掉的也不要算(普通求重心可以不用加vis[v])
		getrt(v, u);
		siz[u] += siz[v];
		maxson[u] = max(maxson[u], siz[v]); //更新maxson
	}
    maxson[u] = max(maxson[u], all - siz[u]);   //考虑u的祖先结点
    if((maxson[u] << 1) <= all) rt2 = rt, rt = u;   //更新重心(最大子树大小最小)
}
//TODO 计算各结点与根结点之间的距离并全部记录在tmp里
void getdis(int u, int f) {
    tmp[cnt++] = dis[u];
    for(int i = head[u]; ~i; i = edge[i].next) {
        int v = edge[i].to;
        if(v == f || vis[v]) continue;
        dis[v] = dis[u] + edge[i].w;
        getdis(v, u);
    }
}
//TODO 处理经过根结点的路径
//! 注意judge数组要存放之前子树里存在的路径长度，排除折返路径的可能
void solve(int u) {
    static std::queue<int> que;
    for(int i = head[u]; ~i; i = edge[i].next) {
        int v = edge[i].to;
        if(vis[v]) continue;
        cnt = 0;              //注意置零计数器
        dis[v] = edge[i].w;
        getdis(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) {
            //把存在的单条路径长度标上true，供下个子树用
            que.push(tmp[j]);
            judge[tmp[j]] = true;
        }
    }
    //清空judge数组，不要用memset
    while(!que.empty()) {
        judge[que.front()] = false;
        que.pop();
    }
}
//TODO 分治
void divide(int u) {
    vis[u] = judge[0] = true;   //删除根结点
    solve(u);               //计算经过根结点的路径
    for(int i = head[u]; ~i; i = edge[i].next) {
        //分治剩余部分
        int v = edge[i].to;
        if(vis[v]) continue;
        rt = rt2 = 0;
        maxson[rt] = all = siz[v];  //把重心置为0，并把maxson[0]置为最大值
        getrt(v, 0);
        getrt(rt, 0);        //与主函数相同，第二次更新siz大小
        divide(rt);
    }
}

int main() {
    memset(head, -1, sizeof(head));
    scanf("%d%d", &n, &m);
    for(int i = 1; i < n; ++i) {
        int u, v, w;
        scanf("%d%d%d", &u, &v, &w);
        addedge(u, v, w);
        addedge(v, u, w);
    }
    for(int i = 0; i < m; ++i) scanf("%d", &q[i]);
    rt = rt2 = 0;
    maxson[rt] = all = n; //maxson[0]置为最大值（一开始rt=0）
    getrt(1, 0);     //找重心
    //! 此时siz数组存放的是以1为根时的各树大小，需要以找出的重心为根重算
    getrt(rt, 0);
    divide(rt);     //找好重心就可以开始分治了
    for(int i = 0; i < m; i++) {
        if(ans[i]) puts("AYE");
        else puts("NAY");
    }
    return 0;
}