P3806 点分治 dsu on tree做法

· · 题解

感言

刚学dsu on tree（树上启发式合并），后来发现这题可以使用树上启发式合并。

调了3天终于过了。纪念一下哈。

正文

由于是刚学的树上启发式合并，所以会着重讲解思路。

树上启发式合并的思想很简单：对于每个子节点，继承它重儿子的信息，再合并轻儿子的信息。而我们一般用一个全局变量来储存这个信息。

那看到这题，求满足条件的点对。想想全局变量存什么信息？

首先要知道为什么树上启发式合并可以优化时间复杂度，是因为子节点继承重儿子的信息是接近 O(1)（即不用对重儿子的信息更新）的，而重儿子占这个节点的子节点的几乎一半，那么合并的轻儿子的数量大大减少，从而减少每个节点向子节点遍历的次数。所以这个全局信息要能够在 O(1) 时间内继承到父节点。

回顾应用容斥原理的朴素算法，我们发现计数器 cnt 很适合做这个全局变量的。因为 dis[u] 存的是 u 到根节点的距离，这个信息在重儿子回到父节点的时候是不会变的，cnt 自然也不会。

那么算法的流程十分清晰了：

  1. 先统计轻儿子的答案，不保留，清空 cnt。
  2. 统计重儿子的答案，保留 cnt。
  3. 合并轻儿子，并在此过程中计算出答案。

详细的点在代码中会提及：

# include <bits/stdc++.h>
# include <iostream>

using namespace std;

// 建图部分

struct Node
{
    int p;
    int v;
    struct Node* nxt;
};
struct Head
{
    struct Node* nxt;
};

struct Head p[10005];

struct Node* ini()
{
    struct Node* tmp = (struct Node*) malloc (sizeof(struct Node));
    tmp->nxt = NULL;
    return tmp;
}

void add(int x,int y,int v)
{
    struct Node* tmp1 = ini();
    struct Node* tmp2 = ini();

    tmp1->v = tmp2->v = v;

    tmp1->nxt = p[x].nxt;
    tmp1->p = y;
    p[x].nxt = tmp1;

    tmp2->nxt = p[y].nxt;
    tmp2->p = x;
    p[y].nxt = tmp2;

    return ;
}

//dus on tree部分

int n,m;
int tot[10005],son[10005];
int dis[10005];//所有点的dis值
int q[10005],cnt[100000005]; // q记录重儿子子树上的所有节点 cnt是dis计数器
int nq[10005]，ncnt[100000005]; //记录子树上所有节点，最后要合并到q中，作为新的子树信息
// ncnt记录链上节点的dis值
int ask[10005];
int ans[10005];
int top,ntop,root; // top，ntop分别处理q和nq

void dfs(int u,int fa)//找重儿子
{
    tot[u] = 1;
    int maxson=0;
    for (struct Node* tmp = p[u].nxt;tmp!=NULL;tmp = tmp->nxt)
    {
        int v = tmp->p;
        if (v == fa)
        {
            continue;
        }
        dfs(v,u);
        tot[u] += tot[v];
        if (tot[v] > maxson)
        {
            maxson = tot[v];
            son[u] = v;
        }
    }
    return ;
}

void merge(int u,int fa)
{
    ncnt[dis[u]]++;//记录u的dis值，计算要用

    for (struct Node* tmp = p[u].nxt;tmp!=NULL;tmp = tmp->nxt)
    {
        int v = tmp->p;
        if (v == fa || v == son[root])
        {
            continue;
        }
        nq[++ntop] = v;//不放在外面的原因是不让root被记录下来
        merge(v,u);
        if (u == root)//若子树遍历完了，要遍历下一棵子树，清空这一棵子树的信息，同时合并这棵子树的信息到q和cnt中
        {
            while (ntop > 0)
            {
                q[++top] = nq[ntop];
                cnt[dis[nq[ntop]]]++;
                ntop--;
            }
        }
    }

    ncnt[dis[u]]--; //不用这个信息了

    for (int i=1;i<=m;i++)
    {
        if (ask[i]+2*dis[root]-dis[u] >= 0 && cnt[ask[i]+2*dis[root]-dis[u]])//先处理不在同一子树上的情况，用的是合并过的cnt
        {
            ans[i] = 1;
        }
        if (dis[u]-ask[i] >= 0 && ncnt[dis[u]-ask[i]])//处理在同一子树上的情况，用的是当前子树的ncnt
        {
            ans[i] = 1;
        }
    }

    return ;
}

void dsu(int u,int fa,bool keep) // dsu on tree主体部分,u为当前节点,fa为父节点,keep为是否保留信息
{
    for (struct Node* tmp = p[u].nxt;tmp!=NULL;tmp = tmp->nxt) // 遍历轻儿子
    {
        int v = tmp->p;
        if (v == fa) // 跳过父节点
        {
            continue;
        }
        dis[v] = dis[u] + tmp->v; // 更新儿子dis值
        if (v == son[u])    // 跳过重儿子
        {
            continue;
        }
        dsu(v,u,false);
    }

    if (son[u]) // 若存在轻儿子
    {
        dsu(son[u],u,true);
        for (int i=1;i<=m;i++) //注意处理重儿子子树内点对，因为是递归调用，所以回到这一层时只用考虑u->son[u]这一条边有无解
        {
            if (cnt[ask[i]+dis[u]])
            {
                ans[i]=1;
            }
        }
    }

    root = u;    // 防止遍历到父节点
    merge(u,fa); // 合并轻儿子

    q[++top] = u;
    cnt[dis[u]]++;

    if (!keep) // 若u是轻儿子
    {
        while (top > 0)//不保留信息，清空
        {
            cnt[dis[q[top]]]--;
            top--;
        }
    }

    return ;
}

int main (void)
{
    scanf ("%d %d",&n,&m);

    for (int i=0;i<n-1;i++)
    {
        int x,y,v;
        scanf ("%d %d %d",&x,&y,&v);
        add(x,y,v);
    }
    for (int i=1;i<=m;i++)
    {
        scanf ("%d",&ask[i]);
    }

    dfs(1,0);
    dsu(1,0,true);

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

    return 0;
}