BZOJ 1316: 树上的询问

挺裸的点分治

刚开始想用map水过去，然后做p次点分治，然后T到自闭

最后发现可以sort一遍，然后去重，记录每个数出现的次数，这样就可以双指针，不会漏掉了

#include <bits/stdc++.h>
using namespace std;
 
#define N 100010
int n, q, Q[N], ans[N]; 
struct Graph
{
    struct node
    {
        int to, nx, w;
        node() {}
        node (int to, int nx, int w) : to(to), nx(nx), w(w) {}
    }a[N << 1];
    int head[N], pos;
    void add(int u, int v, int w)
    {
        a[++pos] = node(v, head[u], w); head[u] = pos; 
        a[++pos] = node(u, head[v], w); head[v] = pos; 
    }
}G;
#define erp(u) for (int it = G.head[u], v = G.a[it].to, w = G.a[it].w; it; it = G.a[it].nx, v = G.a[it].to, w = G.a[it].w)
 
bool vis[N];
int root, sum, sze[N], f[N]; 
void getroot(int u, int fa)
{
    sze[u] = 1, f[u] = 0;
    erp(u) if (v != fa && !vis[v])
    {
        getroot(v, u);
        sze[u] += sze[v];
        f[u] = max(f[u], sze[v]);   
    }
    f[u] = max(f[u], sum - sze[u]);
    if (f[u] < f[root]) root = u; 
}
 
int deep[N], d[N];
void getdeep(int u, int fa)
{
    deep[++deep[0]] = d[u];
    erp(u) if (v != fa && !vis[v])
    {
        d[v] = d[u] + w; 
        getdeep(v, u);
    }
}
 
int g[N], num[10000010]; 
void calc(int u, int cost, int opt)
{
    d[u] = cost; deep[0] = 0;
    getdeep(u, 0);
    sort(deep + 1, deep + 1 + deep[0]);
    g[0] = 0;
    for (int i = 1; i <= deep[0]; ++i)
    {
        if (i == 1 || deep[i] != deep[i - 1]) g[++g[0]] = deep[i], num[g[0]] = 1;
        else ++num[g[0]]; 
    }
    for (int i = 1; i <= q; ++i)
    {
        for (int j = 1; j <= g[0]; ++j) 
        {
            if (num[j] > 1 && 2 * g[j] == Q[i]) 
                ans[i] += opt * num[j] * (num[j] - 1);
        }
        int r = g[0];
        for (int l = 1; l < r; ++l)
        {
            while (l < r && g[l] + g[r] >= Q[i]) 
            {
                if (g[l] + g[r] == Q[i]) ans[i] += opt * num[l] * num[r];
                --r;
            }
        }
    }
}
 
void solve(int u)
{
    calc(u, 0, 1);
    vis[u] = 1;
    erp(u) if (!vis[v]) 
    {
        calc(v, w, -1);
        sum = f[0] = sze[v]; 
        root = 0;
        getroot(v, u); 
        solve(root); 
    }
}
 
void Run()
{
    scanf("%d%d", &n, &q);
    {
        for (int i = 1, u, v, w; i < n; ++i)
        {
            scanf("%d%d%d", &u, &v, &w);
            G.add(u, v, w);
        }
        for (int qq = 1; qq <= q; ++qq) scanf("%d", Q + qq);
        sum = n; f[0] = n;
        root = 0; 
        getroot(1, 0);
        solve(root);
        for (int i = 1; i <= q; ++i) puts(ans[i] || !Q[i] ? "Yes" : "No");
    }
}
 
int main()
{
    Run();
}