BZOJ 3697 树分治

Description

采药人的药田是一个树状结构，每条路径上都种植着同种药材。
采药人以自己对药材独到的见解，对每种药材进行了分类。大致分为两类，一种是阴性的，一种是阳性的。

采药人每天都要进行采药活动。他选择的路径是很有讲究的，他认为阴阳平衡是很重要的，所以他走的一定是两种药材数目相等的路径。采药工作是很辛苦的，所以他希望他选出的路径中有一个可以作为休息站的节点（不包括起点和终点），满足起点到休息站和休息站到终点的路径也是阴阳平衡的。他想知道他一共可以选择多少种不同的路径。

如果只是求路径、满足路径上面两种药材数量相同的数量，直接点分治就可以了。

但是这里还增加了休息站这个条件~~~~~~虐蒟蒻啊~~~~~

由于休息站要么在起点到根要么在根到终点这两个路径~~~反正就根到某一个点的路径~~

在这里维护一个前缀和，dfs维护路径的前缀和，对于当前点，权值为X,看看前缀和X是否存在，如果存在就意味当前点到根有休息站~~

ps....这里的休息站绝对不会是根哟~~~

设f[i][0/1]:=根到当前子树权值为i且没有/有休息站在当子树的数量

设g[i][0/1]:=根到前几颗子树权值为i且没有/有休息站在子树的数量

那么当前子树的贡献为：

f[0][0]*(g[0][0]-1)+∑f[i][1]*g[-i][1]+f[i][0]*g[-i][1]+f[i][1]*g[-i][0]

f[0][0]*(g[0][0]-1)~~~这个式子意味着以根节点为休息站的路径，初始化g[0][0]=1、空集...-1是为了避免只选了一半路径~~

f[i][1]*g[-i][1]指的是在当前子树里有一个休息站、然后再前几颗子树里也选一些路径、这些路径也包含休息站。同时这两个休息站之间也是平衡的~~

最后两个式子就是只有一个休息站~~

#include <algorithm>
#include <iostream>
#include<string.h>
#include <fstream>
#include <math.h>
#include <vector>
#include <cstdio>
#include <string>
#include <queue>
#include <stack>
#include <map>
#include <set>
#define exp 1e-8
#define fi first
#define se second
#define ll long long
#define INF 0x3f3f3f3f
#define pb(a) push_back(a)
#define mp(a,b) make_pair(a,b)
#define all(a) a.begin(),a.end()
#define mm(a,b) memset(a,b,sizeof(a));
#define for0(a,b) for(int a=0;a<=b;a++)//0---(b-1)
#define for1(a,b) for(int a=1;a<=b;a++)//1---(b)
#define rep(a,b,c) for(int a=b;a<=c;a++)//b---c
#define repp(a,b,c)for(int a=b;a>=c;a--)///
#define cnt_one(i) __builtin_popcount(i)
#define stl(c,itr) for(__typeof((c).begin()) itr=(c).begin();itr!=(c).end();itr++)
using namespace std;
void bug(string m="here"){cout<<m<<endl;}
template<typename __ll> inline void READ(__ll &m){__ll x=0,f=1;char ch=getchar();while(!(ch>='0'&&ch<='9')){if(ch=='-')f=-1;ch=getchar();}while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}m=x*f;}
template<typename __ll>inline void read(__ll &m){READ(m);}
template<typename __ll>inline void read(__ll &m,__ll &a){READ(m);READ(a);}
template<typename __ll>inline void read(__ll &m,__ll &a,__ll &b){READ(m);READ(a);READ(b);}
template<typename __ll>inline void read(__ll &m,__ll &a,__ll &b,__ll &c){READ(m);READ(a);READ(b);READ(c);}
template<typename __ll>inline void read(__ll &m,__ll &a,__ll &b,__ll &c,__ll &d){READ(m);READ(a);READ(b);READ(c);read(d);}
template < class T > inline  void out(T a){if(a>9)out(a/10);putchar(a%10+'0');}
template < class T > inline  void outln(T a){if(a>9)out(a/10);putchar(a%10+'0');puts("");}
template < class T > inline  void out(T a,T b){out(a);putchar(' ');out(b);}
template < class T > inline  void outln(T a,T b){out(a);putchar(' ');outln(b);}
template < class T > inline  void out(T a,T b,T c){out(a);putchar(' ');out(b);putchar(' ');out(c);}
template < class T > inline  void outln(T a,T b,T c){out(a);putchar(' ');outln(b);putchar(' ');outln(b);}
template < class T > T gcd(T a, T b) { return b ? gcd(b, a % b) : a; }
template < class T > T lcm(T a, T b) { return a / gcd(a, b) * b; }
template < class T > inline void rmin(T &a, const T &b) { if(a > b) a = b; }
template < class T > inline void rmax(T &a, const T &b) { if(a < b) a = b; }
template < class T > T pow(T a, T b) { T r = 1; while(b > 0) { if(b & 1) r = r * a; a = a * a; b /= 2; } return r; }
template < class T > T pow(T a, T b, T mod) { T r = 1; while(b > 0) { if(b & 1) r = r * a % mod; a = a * a % mod; b /= 2; } return r; }

const int cnt_edge=200100;  //修改啊
const int cnt_v=100100;
int head[cnt_v],cnt_e;
struct EDGE{int u,v,next,cost;}edge[cnt_edge];
void init(){cnt_e=0;memset(head,-1,sizeof(head));}
void addedge(int u,int v,int cost=0)
{edge[cnt_e].u=u;edge[cnt_e].v=v;edge[cnt_e].cost=cost;edge[cnt_e].next=head[u];head[u]=cnt_e++;}

int n,k;
int root,tot;
int minn,size[100100];///minn：衡量某个节点是否能当重心，minn越小 就越可以当重心、size[i] i子树的大小
bool del[100100];///记录是否已经删除这个点...
ll ans,f[200100][2],g[200100][2];
int pre[200100],dist[100100];

void getroot(int u,int fa)
{
    size[u]=1;
    int maxn=0;
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==fa||del[v])continue;///这里必须要加上del[v]
        getroot(v,u);
        size[u]+=size[v];
        maxn=max(maxn,size[v]);
    }
    maxn=max(maxn,tot-size[u]);
    if(maxn<minn)root=u,minn=maxn;
}
int dfs(int u,int fa,int deep)
{
    int max_deep=deep;
    if(pre[dist[u]])f[dist[u]][1]++;
    else f[dist[u]][0]++;
    pre[dist[u]]++;
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==fa||del[v])continue;
        dist[v]=dist[u]+edge[i].cost;
        max_deep=max(max_deep,dfs(v,u,deep+1));
    }
    pre[dist[u]]--;
    return max_deep;
}
void work(int u)
{
    int max_deep=0;
    del[u]=g[n][0]=1;///初始化g[0][0]=1...空集
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(del[v])continue;
        dist[v]=n+edge[i].cost;
        int cur_deep=dfs(v,u,1);
        max_deep=max(max_deep,cur_deep);
        ans+=(g[n][0]-1)*f[n][0];///以根为休息点
        for(int j=-cur_deep;j<=cur_deep;j++)
            ans+=g[n-j][1]*f[n+j][1]+g[n-j][1]*f[n+j][0]+g[n-j][0]*f[n+j][1];
        for(int j=-cur_deep;j<=cur_deep;j++)
        {
            g[n+j][0]+=f[n+j][0];
            g[n+j][1]+=f[n+j][1];
            f[n+j][1]=f[n+j][0]=0;
        }
    }
    for(int j=-max_deep;j<=max_deep;j++)
        g[n+j][1]=g[n+j][0]=0;
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(del[v])continue;
        minn=INF,tot=size[v];
        getroot(v,u);
        work(root);
    }
}
int main()
{
    read(n);
    init();
    for1(i,n-1)
    {
        int a,b,c;
        read(a,b,c);
        if(!c)c=-1;
        addedge(a,b,c);
        addedge(b,a,c);
    }                    ///读入

    ans=0;
    memset(del,0,sizeof del); ///做题初始化

    minn=INF,tot=n;
    getroot(1,0);      ///找到一棵树的重，tot为该树的大小....
    work(root);
    outln(ans);
    return 0;
}