hihocoder 树的黑白染色

http://hihocoder.com/problemset/problem/1823

给你一棵树，然后有两种操作。

每个节点可以是黑色或者白色，初始时所有节点都是白色。

现在希望你依次进行M个操作，每种操作是以下两种之一：

1. 将x节点涂黑

2. 输出x节点到所有黑色节点的距离之和

hihocoder上出现次数很多的一个题。10^5的查询次数来说的话，n*log(n)在python下也能通过。

扫描二维码关注公众号，回复： 3203605 查看本文章

本来想的方法是对于查询与染色，某个为o(1)另一个是o(n)，这样妥妥不行。

对于n*log(n)来说，就是说其中的最大操作复杂度不能超过log(n)，就是树的高度呗。

思路如下：

每个点维持当前点到其所有子节点的染色点距离和d，以及该染色点数量n（包括该点本身）。这样做一个后序遍历是可以通过子节点的两个数值得到根结点的两个值的，这是前提知识。注意d和n的定义，仅关心该点为根的子树的情况。

对于染色操作，因为每个节点的d和n仅仅由其子孙后代决定，所以只需向上更新到根结点即可。

对于查询操作，你可以由父节点的d个n以及本身节点的d和n来推测出本身节点到兄弟姐妹节点的距离信息，假设本身是i，父节点是f，两者距离是dis，那么d[f]-d[i]-n[i]是父节点刨除节点i子树后的d，然后这个距离加上(n[f]-n[i])*dis就是兄弟姐妹到该点距离了。

code:

n,m=[int(x)for x in raw_input().split()]
f=[-1]+[int(x)for x in raw_input().split()]
df=[-1]+[int(x)for x in raw_input().split()]
cl=[0]*n
ns=[0]*n
ds=[0]*n

def add(i):
    if cl[i]!=0:
        return
    cl[i]=1
    global ns,ds
    cnt=0
    while i!=-1:
        ns[i]+=1
        ds[i]+=cnt
        cnt+=df[i]
        i=f[i]

def check(i):
    global ns,ds
    d=ds[i]
    ff=f[i]
    cnt=df[i]
    while ff!=-1:
        node_other=ns[ff]-ns[i]
        dis_other=ds[ff]-ds[i]-ns[i]*df[i]
        dis=dis_other+node_other*cnt
        d+=dis

        i=ff
        ff=f[i]
        cnt+=df[i]
    print(d)
    return d

for _ in range(m):
    t,p=[int(x) for x in raw_input().split()]
    if t==1:
        add(p)
    else:
        check(p)