洛谷P4930 [FJ2014集训]采药人的路径【点分治】

时空限制 1000ms / 128MB

题目描述

采药人的药田是一个树状结构，每条路径上都种植着同种药材。
采药人以自己对药材独到的见解，对每种药材进行了分类。大致分为两类，一种是阴性的，一种是阳性的。
采药人每天都要进行采药活动。他选择的路径是很有讲究的，他认为阴阳平衡是很重要的，所以他走的一定是两种药材数目相等的路径。采药工作是很辛苦的，所以他希望他选出的路径中有一个可以作为休息站的节点（不包括起点和终点），满足起点到休息站和休息站到终点的路径也是阴阳平衡的。他想知道他一共可以选择多少种不同的路径。

输入格式：

第1行包含一个整数N。
接下来N-1行，每行包含三个整数$a_i,b_i$和$t_i$，表示这条路上药材的类型。

输出格式：

输出符合采药人要求的路径数目。

说明

对于100%的数据，N ≤ 100,000

---

题目分析

首先看到这种类型的题肯定 点分治 没跑了

假如我们把给定类型0的边设边权-1，类型1的边设边权1
那么阴阳平衡的路径总长一定为0

假设当前分治的子树根为$rt$，我们要统计该子树内经过$rt$的路径有多少条符合条件
从$rt$开始向每一个子树内dfs，记录下每个结点到$rt$的$dis$

(以下记结点$v$到$rt$的距离为$dis[v]$)
假如某个结点$v_i$有$dis[v_i]=len$，且$v_i$到$rt$的路径上存在另一个结点$v_j$满足$dis[v_j]=len$，就给$v_i$打上标记

于是满足条件的配对就分为4种情况

1. 若结点$v_i$有$dis[v_i]=0$且没有标记，则$v_i$可以与任意满足$dis[v_j]=0$的结点$v_j$配对
2. 若结点$v_i$有$dis[v_i]=0$且有标记，则$v_i$可以与任意满足$dis[v_j]=0$的结点$v_j$配对，以及可以与$rt$配对
3. 若结点$v_i$有$dis[v_i]=len$且没有标记，则$v_i$可以与任意满足$dis[v_j]=-len$且有标记的结点$v_j$配对
4. 若结点$v_i$有$dis[v_i]=len$且有标记，则$v_i$可以与任意满足$dis[v_j]=-len$的结点$v_j$配对

具体实现看代码

---

#include<iostream>
#include<vector>
#include<algorithm>
#include<queue>
#include<cstring>
#include<cstdio>
using namespace std;
typedef long long lt;

int read()
{
    int f=1,x=0;
    char ss=getchar();
    while(ss<'0'||ss>'9'){if(ss=='-')f=-1;ss=getchar();}
    while(ss>='0'&&ss<='9'){x=x*10+ss-'0';ss=getchar();}
    return f*x;
}

const int inf=1e9;
const int maxn=200010;
int n;
struct node{int v,dis,nxt;}E[maxn<<1];
int head[maxn],tot;
int maxp[maxn],size[maxn];
int vis[maxn],sum,rt;
int judge[maxn],q[maxn],p;
int rem[maxn][2],num[maxn][2],cnt0,cnt1;
lt 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)
{
    size[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);
        size[u]+=size[v];
        maxp[u]=max(maxp[u],size[v]);
    }
    maxp[u]=max(maxp[u],sum-size[u]);
    if(maxp[u]<maxp[rt]) rt=u;
}

void getdis(int u,int fa,int dis)
{
    if(judge[dis+n]) rem[++cnt1][1]=dis+n;
    else rem[++cnt0][0]=dis+n;
    //rem[][0/1]存当前子树内所有节点到rt的距离，0表示没有标记，1表示有
    
    judge[dis+n]++;//judge标记出现过的距离
    for(int i=head[u];i;i=E[i].nxt)
    {
        int v=E[i].v;
        if(v==fa||vis[v])continue;
        getdis(v,u,dis+E[i].dis);
    }
    judge[dis+n]--;
}

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;
        
        cnt0=cnt1=0;
        getdis(v,u,E[i].dis);
        
        for(int j=1;j<=cnt0;++j)//分别对应4种配对情况
        {
            int tt=rem[j][0];
            if(tt==n) ans+=num[tt][0]+num[tt][1];
            else ans+=num[-tt+2*n][1];
        }
        for(int j=1;j<=cnt1;++j)
        {
            int tt=rem[j][1];
            if(tt==n) ans+=num[tt][0]+num[tt][1]+1;
            else ans+=num[-tt+2*n][1]+num[-tt+2*n][0];
        }
        
        for(int j=1;j<=cnt0;++j)//num记录前i棵rt的子树距离的出现情况以及是否标记
        num[rem[j][0]][0]++,q[++p]=rem[j][0];
        for(int j=1;j<=cnt1;++j)
        num[rem[j][1]][1]++,q[++p]=rem[j][1];
    }
    
    for(int i=1;i<=p;++i)//清空num数组不要memset，会T
    num[q[i]][0]=num[q[i]][1]=0;
}

void solve(int u)
{   
    vis[u]=1; calc(u);
    for(int i=head[u];i;i=E[i].nxt)
    {
        int v=E[i].v;
        if(vis[v])continue;
        sum=size[v]; maxp[rt=0]=inf;
        getrt(v,0); solve(rt);
    }
}

int main()
{
    n=read();
    for(int i=1;i<n;++i)
    {
        int u=read(),v=read(),dis=read();
        if(dis==0) dis=-1;
        add(u,v,dis);add(v,u,dis);
    }

    maxp[rt]=sum=n;
    getrt(1,0); 
    solve(rt);
    
    printf("%lld",ans);
    return 0;
}