luogu4930

P4930 「FJ2014集训」采药人的路径

题目描述

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

输入输出格式

输入格式：

 

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

 

输出格式：

 

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

输入输出样例

输入样例#1： 复制

7
1 2 0
3 1 1
2 4 0
5 2 0
6 3 1
5 7 1

输出样例#1： 复制

1

说明

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

sol：点分治（板子题），把0看成-1，没有条件2的限制就是求和是0的路径条数，用dp[i]表示和为i的方案数，有了2的限制就要多一维记录当前这条路径是否可以有断点，断点就是记录当前的dis前面是否出现过，即是否当前路径有一段后缀为0

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
inline ll read()
{
  ll s=0; bool f=0; char ch=' ';
  while(!isdigit(ch)) {f|=(ch=='-'); ch=getchar();}
  while(isdigit(ch)) {s=(s<<3)+(s<<1)+(ch^48); ch=getchar();}
  return (f)?(-s):(s);
}
#define R(x) x=read()
inline void write(ll x)
{
  if(x<0) {putchar('-'); x=-x;}
  if(x<10) {putchar(x+'0'); return;}
  write(x/10); putchar((x%10)+'0');
}
#define W(x) write(x),putchar(' ')
#define Wl(x) write(x),putchar('
')
const int N=200005,M=400005,inf=0x3f3f3f3f;
int n;
ll ans=0;
int tot=0,Next[M],to[M],val[M],head[N];
int rt,Sum,mx[N],sz[N];
int dis[N],dep[N],mxdep,arrdis[M];
bool Vis[N];
int dp[2][2][M];
inline void Link(int x,int y,int z)
{
  Next[++tot]=head[x]; to[tot]=y; val[tot]=z; head[x]=tot;
}
inline void getrt(int x,int fat)
{
  int e;
  sz[x]=1; mx[x]=0;
  for(e=head[x];e;e=Next[e]) if((to[e]!=fat)&&(!Vis[to[e]]))
  {
    getrt(to[e],x); sz[x]+=sz[to[e]]; mx[x]=max(mx[x],sz[to[e]]);
  }
  mx[x]=max(mx[x],Sum-sz[x]);  if(mx[x]<mx[rt]) rt=x;
}
inline void calc(int x,int fat,int o)
{
  int e;
  mxdep=max(mxdep,dep[x]);
  if(arrdis[dis[x]+n]) dp[o][1][dis[x]+n]++;
  else dp[o][0][dis[x]+n]++;
  arrdis[dis[x]+n]++;
  for(e=head[x];e;e=Next[e]) if((!Vis[to[e]])&&(to[e]!=fat))
  {
    dep[to[e]]=dep[x]+1; dis[to[e]]=dis[x]+val[e]; calc(to[e],x,o);
  }
  arrdis[dis[x]+n]--;
}
inline void dfs(int x)
{
  int e,i,mm=0;
  Vis[x]=1; dp[0][0][n]=1;
  for(e=head[x];e;e=Next[e]) if(!Vis[to[e]])
  {
    dis[to[e]]=val[e]; dep[to[e]]=1;
        mxdep=1;
    calc(to[e],x,1);
        mm=max(mm,mxdep);
    ans=ans+(dp[0][0][n]-1)*dp[1][0][n];
    for(i=-mxdep;i<=mxdep;i++)
    {
      ans=ans+dp[0][1][n+i]*dp[1][1][n-i]+dp[0][0][n+i]*dp[1][1][n-i]+dp[0][1][n+i]*dp[1][0][n-i];
    }
    for(i=n-mxdep;i<=n+mxdep;i++)
    {
      dp[0][0][i]+=dp[1][0][i];
      dp[0][1][i]+=dp[1][1][i];
      dp[1][0][i]=dp[1][1][i]=0;
    }
  }
  // puts("#####################################");
  for(i=n-mm;i<=n+mm;i++) dp[0][0][i]=dp[0][1][i]=0;
  for(e=head[x];e;e=Next[e]) if(!Vis[to[e]])
  {
    Sum=sz[to[e]]; mx[rt=0]=inf; getrt(to[e],x); dfs(rt);
  }
}
int main()
{
  int i,x,y,z;
  R(n);
  for(i=1;i<n;i++)
  {
    R(x); R(y); R(z); if(z==0) z=-1;
    Link(x,y,z); Link(y,x,z);
  }
  Sum=n; mx[rt=0]=inf;
  getrt(1,0);
  // cout<<"rt="<<rt<<endl;
  dfs(rt);
  Wl(ans);
  return 0;
}