【P3806】點分治模闆

#include<bits/stdc++.h>
using namespace std;
const int inf = 10000000;
const int maxn = 100010;
int n, m;
struct node { int v, dis, nxt; }E[maxn << 1];
int tot, head[maxn];
int maxp[maxn], siz[maxn], dis[maxn], rem[maxn];
int vis[maxn], test[inf], judge[inf], q[maxn];
int query[1010];
int sum, rt;
int ans;

void add(int u, int v, int dis)
{
  E[++tot].nxt = head[u];
  E[tot].v = v;
  E[tot].dis = dis;
  head[u] = tot;
}

void getrt(int u, int pa)
{
  siz[u] = 1;
  maxp[u] = 0;
  for (int i = head[u]; i; i = E[i].nxt)
  {
    int v = E[i].v;
    if (v == pa || vis[v])continue;
    getrt(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 fa)
{
  rem[++rem[0]] = dis[u];
  for (int i = head[u]; i; i = E[i].nxt)
  {
    int v = E[i].v;
    if (v == fa || vis[v])continue;
    dis[v] = dis[u] + E[i].dis;
    getdis(v, u);
  }
}


void calc(int u)
{
  int p = 0;
  for (int i = head[u]; i; i = E[i].nxt)
  {
    int v = E[i].v;
    if (vis[v])continue;
    rem[0] = 0; 
    dis[v] = E[i].dis;
    getdis(v, u);

    for (int j = rem[0]; j; --j)
      for (int k = 1; k <= m; ++k)
        if (query[k] >= rem[j])
          test[k] |= judge[query[k] - rem[j]];

    for (int j = rem[0]; j; --j)
      q[++p] = rem[j], judge[rem[j]] = 1;
  }
  for (int i = 1; i <= p; ++i)
    judge[q[i]] = 0;
}

void solve(int u)
{
  vis[u] = judge[0] = 1; 
  calc(u);
  for (int i = head[u]; i; i = E[i].nxt)//對每個子樹進行分治
  {
    int v = E[i].v;
    if (vis[v])continue;
    sum = siz[v];
    maxp[rt = 0] = inf;
    getrt(v, 0); 
    solve(rt);//在子樹中找重心并遞歸處理
  }
}

int main()
{
  scanf("%d%d", &n, &m);
  for (int i = 1; i < n; ++i)
  {
    int u, v, dis;
    scanf("%d%d%d", &u, &v, &dis);
    add(u, v, dis); 
    add(v, u, dis);
  }
  for (int i = 1; i <= m; ++i)
    scanf("%d", &query[i]);

  maxp[rt] = sum = n;
  getrt(1, 0);
  solve(rt);

  for (int i = 1; i <= m; ++i)
  {
    if (test[i]) printf("AYE\n");
    else printf("NAY\n");
  }
  return 0;
}