USACO 2013open ：yinyang（treap+启发式合并）

题目大意：给出一颗树，由N个结点，编号1至N。有N-1条边，每条边都有1只奶牛，奶牛要么是黑毛奶牛，要么是白毛奶牛。从结点A走到结点B,如果不能走重复的边，那么这条路径是唯一的，这样的路径在图论里被称为“简单路径”。我们用（A,B）来表示从A走到B的这条简单路径。如果简单路径（A,B）同时满足以下3个条件，那么简单路径（A,B）被称为“平衡”的简单路径：
至少存在一个结点C，结点C同时满足：
1、结点C是简单路径（A，B）上的一个结点，且结点C不是结点A，结点C也不是结点B。
2、从结点A走到结点C，途中遇到的白毛奶牛的数量等于遇到的黑毛奶牛的数量。
3、从结点C走到结点B，途中遇到的白毛奶牛的数量等于遇到的黑毛奶牛的数量。
你的任务是计算：总共有多少条“平衡”的简单路径？
需要特别注意的有：
1、简单路径（A,B）和简单路径（B，A）被认为是同一条简单路径。
2、如果简单路径（A,B）是“平衡”的，有可能该路径上有多个不同的结点C使得简单路径（A，B）是“平衡”的，但简单路径（A,B）只能算一条“平衡”的，即答案只加1。

Input：

第一行，一个整数N。1 <= N <= 100000。
接下来有N行，每行三个整数： a, b, c。表示结点a与结点b有一条边，如果c=0表示该边上有一只白毛奶牛，c=1表示该边上有一只黑毛奶牛。

Output:

平衡路径的条数。

样例：

Input：

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

Output：

1
【样例解释】
“平衡”的简单路径：3-1-2-5-7，其中2号结点可以被认为是结点C。

题目分析：本题可以用Treap+启发式合并做。白毛奶牛认为其边权为1，黑则为-1。我们用Treap[node]记录以node为根的子树的所有dep，并且记录个数。对于每一个treap上的节点v，用cnt[0]记录当前子树中dep[son]==v，且son的祖先中没有dep[fa]==v的son的个数，用cnt[1]记录祖先中有dep[fa]==v的son的个数。由于一条路径权值和==0，故dep[v]==dep[lca]*2-dep[u]，我们要用son合并node的时候（假设size[son]<size[node]），先枚举treap[son]的节点P，然后在treap[node]中找值为2*dep[node]-P->val的节点Q。由于路径中间必须要有一个中间节点c，故不能用P->cnt[0]*Q->cnt[0]更新答案，其余的两两互乘即可。做完之后记得更新一下dep[node]的那个节点，时间复杂度O(n*log^2(n))。

    之前代码一直WA不知为什么，后来才知道我有一个语句“if (!Q) return;”放错位置了，return了之后就没有递归P的左右节点，导致ans变小，应该先递归。

CODE：

#include<iostream>
#include<string>
#include<cstring>
#include<cmath>
#include<cstdio>
#include<cstdlib>
#include<stdio.h>
#include<algorithm>
#include<ctime>
using namespace std;

const int maxn=100100;
const int maxl=30;

struct Tnode
{
int val,fix;
int cnt[2],Size;
Tnode *lson,*rson;

int Left_size()
{
if (lson) return lson->Size;
return 0;
}

int Right_size()
{
if (rson) return rson->Size;
return 0;
}

void Get_size()
{
Size=Left_size()+Right_size()+1;
}
} tree[maxn*maxl];
Tnode *Root[maxn];
int Tcur=-1;

struct data
{
int obj,len;
data *Next;
} e[maxn<<1];
data *head[maxn];
int Ecur=-1;

int dep[maxn];
int n;
long long ans=0;

void Add(int x,int y,int v)
{
Ecur++;
e[Ecur].obj=y;
e[Ecur].len=v;
e[Ecur].Next=head[x];
head[x]=e+Ecur;
}

Tnode *New_node(int v,int num0,int num1)
{
Tcur++;
tree[Tcur].val=v;
tree[Tcur].fix=rand();
tree[Tcur].lson=tree[Tcur].rson=NULL;
tree[Tcur].cnt[0]=num0;
tree[Tcur].cnt[1]=num1;
tree[Tcur].Size=1;
return tree+Tcur;
}

void Right_turn(Tnode *&P)
{
Tnode *W=P->lson;
P->lson=W->rson;
W->rson=P;
P=W;
P->rson->Get_size();
P->Get_size();
}

void Left_turn(Tnode *&P)
{
Tnode *W=P->rson;
P->rson=W->lson;
W->lson=P;
P=W;
P->lson->Get_size();
P->Get_size();
}

void Insert(Tnode *&P,int v,int num0,int num1)
{
if (!P) P=New_node(v,num0,num1);
else
{
if ( v==P->val )
{
P->cnt[0]+=num0;
P->cnt[1]+=num1;
}
else
{
if ( v < P->val )
{
Insert(P->lson,v,num0,num1);
if ( P->lson->fix < P->fix ) Right_turn(P);
}
else
{
Insert(P->rson,v,num0,num1);
if ( P->rson->fix < P->fix ) Left_turn(P);
}
P->Get_size();
}
}
}

Tnode *Get(Tnode *&P,int v)
{
if (!P) return NULL;
if ( v==P->val ) return P;
if ( v < P->val ) return Get(P->lson,v);
return Get(P->rson,v);
}

void Work(Tnode *&root,Tnode *&P,int lcadep)
{
if (!P) return;
Work(root,P->lson,lcadep);
Work(root,P->rson,lcadep);
Tnode *Q=Get(root,2*lcadep-P->val);
if (!Q) return;
ans+=( (long long)(P->cnt[0])*(long long)(Q->cnt[1]) );
ans+=( (long long)(P->cnt[1])*(long long)(Q->cnt[0]) );
ans+=( (long long)(P->cnt[1])*(long long)(Q->cnt[1]) );
//printf("%d %d\n",P->val,Q->val);
//printf("%d %d\n",P->cnt[0],P->cnt[1]);
//printf("%d %d\n\n",Q->cnt[0],Q->cnt[1]);
}

void Copy(Tnode *&root,Tnode *&P)
{
if (!P) return;
Insert(root,P->val,P->cnt[0],P->cnt[1]);
Copy(root,P->lson);
Copy(root,P->rson);
}

void Update(Tnode *&x,Tnode *&y,int lcadep)
{
Work(x,y,lcadep);
Copy(x,y);
}

void Reset(Tnode *&P,int v)
{
if ( v==P->val )
{
P->cnt[1]+=(P->cnt[0]-1);
P->cnt[0]=1;
return;
}
if ( v < P->val ) Reset(P->lson,v);
else Reset(P->rson,v);
}

void Print(Tnode *&P)
{
if (!P) return;
printf("%d %d %d\n",P->val,P->cnt[0],P->cnt[1]);
Print(P->lson);
Print(P->rson);
}

void Dfs(int node,int fa)
{
Root[node]=NULL;
Insert(Root[node],dep[node],1,0);
for (data *p=head[node]; p; p=p->Next)
if (p->obj!=fa)
{
int son=p->obj;
dep[son]=dep[node]+p->len;
Dfs(son,node);
if ( Root[node]->Size > Root[son]->Size ) Update(Root[node],Root[son],dep[node]);
else
{
Update(Root[son],Root[node],dep[node]);
Root[node]=Root[son];
}
Reset(Root[node],dep[node]);
}

//printf("%d\n",node);
//Print(Root[node]);
//printf("\n");
}

int main()
{
freopen("yinyang.in","r",stdin);
freopen("yinyang.out","w",stdout);

srand( time(0) );

scanf("%d",&n);
for (int i=1; i<=n; i++) head[i]=NULL;
for (int i=1; i<n; i++)
{
int a,b,c;
scanf("%d%d%d",&a,&b,&c);
if (!c) c=-1;
Add(a,b,c);
Add(b,a,c);
}

dep[1]=0;
Dfs(1,1);

//for (int i=1; i<=n; i++) printf("%d ",dep[i]);
//printf("\n");

printf("%I64d\n",ans);
return 0;
}