P11295 [NOISG 2022 Qualification] Dragonfly

P11295 [NOISG 2022 Qualification] Dragonfly

题目背景

在植物园和碧山公园周围的池塘中，经常可以看到蜻蜓。在一个更密集的森林区域中，Benson 记录了 \(n\) 个池塘，以及蜻蜓可以捕食的昆虫数量和种类。

在池塘 \(i\)，共有 \(b_i\) 只昆虫，这些昆虫属于种类 \(s_i\)。

此外，还有 \(n-1\) 条小径，每条小径连接两池塘 \(u_j\) 和 \(v_j\)（双向）。并且满足蜻蜓从每一个池塘出发都能到达其它所有池塘。

Benson 抓住了 \(d\) 只蜻蜓，并计划依次释放到池塘 \(1\)。每只蜻蜓有一个目标池塘 \(h_k \neq 1\)，会沿着最短路径飞到目标池塘，并在经过的每个池塘中捕食昆虫（包括池塘 \(1\)）。

每次捕食会减少池塘中 \(1\) 只昆虫（如果昆虫数量不为 \(0\)）。需要帮助 Benson 计算每只蜻蜓的飞行过程中捕食到的不同种类昆虫的数量。

题目描述

请确定每只蜻蜓的飞行过程中捕食到的不同种类昆虫的数量。

输入格式

• 第一行包含两个整数 \(n\) 和 \(d\)，分别表示池塘数量和蜻蜓数量。
• 第二行包含 \(n\) 个整数 \(b_1, b_2, \dots, b_n\)，表示每个池塘的昆虫数量。
• 第三行包含 \(n\) 个整数 \(s_1, s_2, \dots, s_n\)，表示每个池塘昆虫的种类。
• 第四行包含 \(d\) 个整数 \(h_1, h_2, \dots, h_d\)，表示每只蜻蜓的目标池塘。
• 接下来 \(n-1\) 行，每行包含两个整数 \(u_j\) 和 \(v_j\)，表示一条连接池塘 \(u_j\) 和 \(v_j\) 的小径。

输出格式

输出一行包含 \(d\) 个整数，第 \(k\) 个整数表示第 \(k\) 只蜻蜓捕食到的不同种类昆虫数量。

【数据范围】

• \(2 \leq n \leq 2 \times 10^5\)
• \(1 \leq d \leq 2 \times 10^6\)
• \(1 \leq s_i \leq n\)，\(0 \leq b_i \leq d\)，\(1 \leq u_j, v_j < n\)

Solution:

首先我们将询问看做时间轴，每个询问 \(Q_i\) 对应时间 \(tim_i\) 维护一个东西：\(t_x\) 表示节点 \(x\) 上的最后一只昆虫被吃掉的时间。

假设现在我们已经维护出了 \(t_x\) 数组，那么我们在面对询问时，只需要对每个颜色使用 set 来维护该颜色 \(col\) 在这条路径 (1,x) 上的所有出现时间 \(t\) 的最小值。如果这个值 \(t_{latest}\ge tim_i\) 那么这个颜色就是有贡献的。

然后我们发现路径的起点一定是 \(1\) 。所以我们将询问挂到点 \(x\) 上。每次进入点 \(x\) 时使用 \(t_x\) 去更新 \(s_x\) 所对应的 set 。然后维护一个序列 \(a_i\) 用来表示当前状态下 \(t_{latest}=i\) 的数量，那么我们每次要查的答案就是序列 \(a_i\) 在 \([tim_i,n]\) 上的区间和。所以我们使用树状数组来维护序列 \(a_i\) 。支持单点修改，区间查询。

然后补充一下最开头的问题：维护 \(t_x\)。

之前我们说将询问挂到点上，这里我们采用类似的思想，我们发现所有能使得点 \(x\) 的昆虫个数减小的路径的终点一定落在点 \(x\) 本身及其子树内。
我们对于时间序列维护一颗权值线段树，向上做线段树合并将挂在点 \(x\) 上的所有询问插入，对每个点线段树上二分一个点 \(pos\) 。使得区间 \([1,pos]\) 上的点的个数 \(cnt\) 恰好等于 \(b_i\) 若没有，则 \(t_x=n\)

时间复杂度 \(O(n\log n)\)，然后这题就做完了。

Code:

#include<bits/stdc++.h>
#define ll long long
#define Max(x,y) (x > y ? x : y)
#define Min(x,y) (x < y ? x : y)
using namespace std;
const int N=2e6+5;
const int D=2e6+5;
const int lg=20;
const int MX=D*lg;
const int inf=1e9;
int a[N],t[N],b[N],s[N],h[D];
int n,m;
vector<int> E[N];
struct Segment_Tree{
    #define ls(x) t[x].Ls
    #define rs(x) t[x].Rs
    #define mid (l+r>>1)
    struct Tree{
        int Ls,Rs,siz;
    }t[D*lg];int cnt=0;
    inline void pushup(int x)
    {
        t[x].siz=t[ls(x)].siz+t[rs(x)].siz;
    }
    void merge(int &x,int y,int l,int r)
    {
        if(!x||!y){x=x|y;return;}
        if(l==r){t[x].siz+=t[y].siz;return;}
        merge(ls(x),ls(y),l,mid);merge(rs(x),rs(y),mid+1,r);
        pushup(x);
    }
    void upd(int &x,int l,int r,int pos)
    {
        t[x=(x?x:++cnt)].siz++;if(l==r)return;
        if(pos<=mid)upd(ls(x),l,mid,pos);
        if(mid<pos)upd(rs(x),mid+1,r,pos);
    }
    int query(int x,int l,int r,int k)
    {
        if(l==r)return l;
        if(k<=t[ls(x)].siz)return query(ls(x),l,mid,k);
        else return query(rs(x),mid+1,r,k-t[ls(x)].siz);
    }
    #undef ls(x)
    #undef rs(x)
    #undef mid
}T;
int rt[N];
multiset<int> S[N];
vector<int> q[N];
int ans[D];
struct Bit{
    int bit[D];
    int test[D];
    inline int lowbit(int x){return  x & (-x);}
    inline void add(int x,int k)
    {
        while(x<=m){bit[x]+=k;x+=lowbit(x);}
    }
    inline int sum(int x)
    {
        int ans=0;
        while(x)ans+=bit[x],x-=lowbit(x);
        return ans;
    }
    inline int query(int l,int r)
    {
        if(!l)return 0;
        return sum(r)-sum(l-1);
    }
}B;
void dfs(int x,int fa)
{
    for(int y : E[x])if(y!=fa)
    {
        dfs(y,x);
        T.merge(rt[x],rt[y],1,m);
    }
    for(int pos : q[x])T.upd(rt[x],1,m,pos);
    if(b[x])
    {
        if(T.t[rt[x]].siz<b[x])t[x]=m;
        else t[x]=T.query(rt[x],1,m,b[x]);
    }
}
void add(int col,int val)
{
    int last =(*S[col].begin());
    S[col].insert(-val);
    if(S[col].size()==1){B.add(val,1);return;}
    int now=(*S[col].begin());
    if(last!=now){B.add(-last,-1),B.add(-now,1);}
}
void del(int col,int val)
{
    int last =(*S[col].begin());
    S[col].erase(S[col].find(-val));
    if(S[col].empty()){B.add(val,-1);return;}
    int now= (*S[col].begin());
    if(last!=now){B.add(-last,-1),B.add(-now,1);}
}
void calc(int x,int fa)
{
    if(t[x])add(s[x],t[x]);
    for(int id : q[x])ans[id]=B.query(id,m);
    for(int y : E[x])if(y!=fa)calc(y,x);
    if(t[x])del(s[x],t[x]);
}
void work()
{
    cin>>n>>m;
    for(int i=1;i<=n;i++)scanf("%d",&b[i]);
    for(int i=1;i<=n;i++)scanf("%d",&s[i]);
    for(int i=1;i<=m;i++)scanf("%d",&h[i]),q[h[i]].push_back(i);
    for(int i=1,u,v;i<n;i++)
    {
        scanf("%d%d",&u,&v);
        E[u].push_back(v);E[v].push_back(u);
    }
    dfs(1,0);calc(1,0);
    for(int i=1;i<=m;i++)
    {
        printf("%d ",ans[i]);
    }
}
int main()
{
    //freopen("P11295_24.in","r",stdin);
    //freopen("woodpecker.out","w",stdout);
    work();
    return 0;
}