@codechef - TREEPATH@ Decompose the Tree

---

@description@

给定一棵无根树，每个节点上都写了一个整数。
你的任务就是统计有多少种方法可以将这棵树分解为若干条路径，使得每个节点恰好属于一条路径，而且每条路径的节点上的数字之和非负。

输入格式
输入数据第一行包含一个整数 T，表示数据组数。
接下来是 T 组数据。 每组数据的第一行包含一个整数 N，代表树中结点个数。
接下来一行包含 N 个整数，由空格分隔，代表每个节点上写的数字。
接下来 N −1 行，每行包含两个整数 Xj 和 Yj，代表编号为 Xj 和 Yj 的节点之间有边直接相连。

输出格式
对于每组数据，输出一行，包含一个整数，即为将树分解的方案数对 10^9 +7 取模得到的结果。

数据范围与子任务
• 1 ≤ T ≤ 10^5
• 1 ≤ Xj,Yj ≤ N

样例数据
输入
1 4
1 10 5 -1
1 2
1 3
2 4
输出
4

样例解释
一共有 4 种分解方法：
• 整棵树即为一条路径，其和为 1 + 10 + 5 + (−1) = 15；
• 一条路径包含节点 2 和 4，其和为 10 + (−1) = 9；另一条路径包含节点 1 和 3，其和为 1 + 5 = 6；
• 一条路径包含节点1、2和4，其和为 1+10+(−1) = 10；另一条路径包含节点3，其和为 5；
• 第一条路径包含节点 2 和 4，其和为 10 + (−1) = 9；第二条路径包含节点 1，其和为 1；第 三条路径包含节点 3，其和为 3。

@solution@

首先可以想到一个朴素 dp：定义 dp[x] 表示分解 x 为根的子树的合法方案数。
转移时枚举一条以 x 为 lca 的非零链，将链所有支出去的枝丫的 dp 值乘起来即可。
这样子搞是 O(n^2)，考虑优化。

我们考虑可以使用启发式合并 + 平衡树来优化我们枚举链的过程。
考虑处理完轻子树继承重子树，我们维护重子树中每个儿子向上爬到当前点的链权值和以及枝丫的 dp 值乘积。
每次通过打 tag 给链权值和加上当前点的 val，方案数乘上轻子树的 dp 值乘积。

然后考虑处理轻子树。先统计链的一个端点在当前轻子树，另一个端点维护在平衡树里的情况。
这个可以通过平衡树维护子树方案数之和 + 分裂成权值 < k 的与 >= k 的两棵平衡树，然后取 >= k 的平衡树的答案搞定。
然后就把这棵轻子树里所有的链丢进平衡树里。

最后加入当前根，记得再在平衡树取出权值和 >= 0 的链的方案数（即链的端点为当前点的情况）。

看起来非常完美，然而有一个问题：万一 dp 值为 0（可能本身没有方案也可能方案数太多 mod 10^9 + 7 = 0），是不存在逆元的。
当然解决方法很多，但是我选择最简单（最sb）的解决方法：判断它有多少的儿子的 dp 值为 0。
如果多于 2 个显然就全部方案数都为 0，否则分类讨论一大堆，然后 WA 到怀疑人生。。。

@accepted code@

#include<cstdio>
#include<queue>
#include<cstdlib>
#include<iostream>
#include<algorithm>
using namespace std;
const int MAXN = 100000;
const int MOD = int(1E9) + 7;
inline int add(int a, int b) {return (a + b)%MOD;}
inline int mul(int a, int b) {return 1LL*a*b%MOD;}
int pow_mod(int b, int p) {
    int ret = 1;
    while( p ) {
        if( p & 1 ) ret = 1LL*ret*b%MOD;
        b = 1LL*b*b%MOD;
        p >>= 1;
    }
    return ret;
}
struct edge{
    edge *nxt; int to;
}edges[2*MAXN + 5], *adj[MAXN + 5], *ecnt=&edges[0];
void addedge(int u, int v) {
    edge *p = (++ecnt);
    p->to = v, p->nxt = adj[u], adj[u] = p;
    p = (++ecnt);
    p->to = u, p->nxt = adj[v], adj[v] = p;
}
int siz[MAXN + 5], hvy[MAXN + 5];
void dfs1(int x, int f) {
    siz[x] = 1, hvy[x] = 0;
    for(edge *p=adj[x];p;p=p->nxt) {
        if( p->to == f ) continue;
        dfs1(p->to, x);
        siz[x] += siz[p->to];
        if( siz[p->to] > siz[hvy[x]] )
            hvy[x] = p->to;
    }
}
struct Treap{
    struct node{
        node *ch[2];
        int key, cnt, sum;
        int tg1, tg2;
        unsigned int pri;
    }pl[MAXN + 5], *root, *NIL, *ncnt;
    typedef pair<node*, node*> Droot;
    Treap() {
        ncnt = root = NIL = &pl[0];
        NIL->ch[0] = NIL->ch[1] = NIL; NIL->sum = 0;
    }
    unsigned int get_rand() {
        return (rand() << 16) | rand() ^ (rand() << 8);
    }
    node *newnode(int k, int c) {
        node *f = (++ncnt);
        f->ch[0] = f->ch[1] = NIL;
        f->pri = get_rand(), f->key = k, f->cnt = f->sum = c;
        f->tg1 = 0, f->tg2 = 1;
        return f;
    }
    void pushdown(node *x) {
        if( x->tg1 ) {
            if( x->ch[0] != NIL )
                x->ch[0]->tg1 += x->tg1, x->ch[0]->key += x->tg1;
            if( x->ch[1] != NIL )
                x->ch[1]->tg1 += x->tg1, x->ch[1]->key += x->tg1;
            x->tg1 = 0;
        }
        if( x->tg2 != 1 ) {
            if( x->ch[0] != NIL ) {
                x->ch[0]->tg2 = mul(x->ch[0]->tg2, x->tg2);
                x->ch[0]->cnt = mul(x->ch[0]->cnt, x->tg2);
                x->ch[0]->sum = mul(x->ch[0]->sum, x->tg2);
            }
            if( x->ch[1] != NIL ) {
                x->ch[1]->tg2 = mul(x->ch[1]->tg2, x->tg2);
                x->ch[1]->cnt = mul(x->ch[1]->cnt, x->tg2);
                x->ch[1]->sum = mul(x->ch[1]->sum, x->tg2);
            }
            x->tg2 = 1;
        }
    }
    void pushup(node *x) {
        x->sum = add(x->cnt, add(x->ch[0]->sum, x->ch[1]->sum));
    }
    node *merge(node *rt1, node *rt2) {
        if( rt1 == NIL ) return rt2;
        if( rt2 == NIL ) return rt1;
        pushdown(rt1), pushdown(rt2);
        if( rt1->pri < rt2->pri ) {
            rt2->ch[0] = merge(rt1, rt2->ch[0]), pushup(rt2);
            return rt2;
        }
        else {
            rt1->ch[1] = merge(rt1->ch[1], rt2), pushup(rt1);
            return rt1;
        }
    }
    Droot split(node *rt, int k) {
        if( rt == NIL ) return make_pair(NIL, NIL);
        pushdown(rt);
        if( k <= rt->key ) {
            Droot tmp = split(rt->ch[0], k);
            rt->ch[0] = tmp.second; pushup(rt);
            return make_pair(tmp.first, rt);
        }
        else {
            Droot tmp = split(rt->ch[1], k);
            rt->ch[1] = tmp.first; pushup(rt);
            return make_pair(rt, tmp.second);
        }
    }
    void insert(node *x) {
        Droot tmp = split(root, x->key);
        root = merge(merge(tmp.first, x), tmp.second);
    }
    int query(int k) {
        Droot tmp = split(root, k);
        int ret = tmp.second->sum;
        root = merge(tmp.first, tmp.second);
        return ret;
    }
    void update1(int x) {
        if( root != NIL )
            root->tg1 += x, root->key += x;
    }
    void update2(int x) {
        if( root != NIL )
            root->tg2 = mul(root->tg2, x), root->cnt = mul(root->cnt, x), root->sum = mul(root->sum, x);
    }
    void init() {
        ncnt = root = NIL;
    }
}T;
int dp[MAXN + 5], spro[MAXN + 5], inv[MAXN + 5], pos[MAXN + 5], val[MAXN + 5];
int dfs3(int x, int f, int s, int k) {
    int ret = 0;
    if( !pos[x] ) ret = mul(T.query(-s), mul(k, spro[x]));
    for(edge *p=adj[x];p;p=p->nxt) {
        if( p->to == f ) continue;
        if( pos[x] ) {
            if( pos[x] == p->to )
                ret = add(ret, dfs3(p->to, x, s + val[p->to], mul(k, spro[x])));
        }
        else ret = add(ret, dfs3(p->to, x, s + val[p->to], mul(k, mul(spro[x], inv[p->to]))));
    }
    return ret;
}
void dfs4(int x, int f, int s, int k) {
    if( !pos[x] ) T.insert(T.newnode(s, mul(k, spro[x])));
    for(edge *p=adj[x];p;p=p->nxt) {
        if( p->to == f ) continue;
        if( pos[x] ) {
            if( pos[x] == p->to )
                dfs4(p->to, x, s + val[p->to], mul(k, spro[x]));
        }
        else dfs4(p->to, x, s + val[p->to], mul(k, mul(spro[x], inv[p->to])));
    }
}
void update(int x, int y) {
    if( !dp[y] ) pos[x] = (pos[x] ? -1 : y);
    else spro[x] = mul(spro[x], dp[y]);
}
void dfs2(int x, int f, bool flag) {
    spro[x] = 1, pos[x] = 0, dp[x] = 0;
    if( !hvy[x] ) {
        inv[x] = dp[x] = (val[x] >= 0);
        if( flag )
            T.insert(T.newnode(val[x], 1));
        return ;
    }
    for(edge *p=adj[x];p;p=p->nxt) {
        if( p->to == f || p->to == hvy[x] ) continue;
        dfs2(p->to, x, false), update(x, p->to);
    }
    dfs2(hvy[x], x, true), update(x, hvy[x]);
    T.update1(val[x]);
    if( !pos[x] ) {
        T.update2(mul(spro[x], inv[hvy[x]]));
        for(edge *p=adj[x];p;p=p->nxt) {
            if( p->to == f || p->to == hvy[x] ) continue;
            dp[x] = add(dp[x], dfs3(p->to, x, val[p->to], inv[p->to]));
            dfs4(p->to, x, val[x] + val[p->to], mul(spro[x], inv[p->to]));
        }
        T.insert(T.newnode(val[x], spro[x]));
    }
    else {
        if( pos[x] == hvy[x] ) {
            T.update2(spro[x]);
            for(edge *p=adj[x];p;p=p->nxt) {
                if( p->to == f || p->to == hvy[x] ) continue;
                dp[x] = add(dp[x], dfs3(p->to, x, val[p->to], inv[p->to]));
            }
        }
        else {
            if( pos[x] != -1 ) {
                dp[x] = add(dp[x], dfs3(pos[x], x, val[pos[x]], mul(spro[x], inv[hvy[x]])));
                T.update2(0);
                dfs4(pos[x], x, val[x] + val[pos[x]], spro[x]);
                for(edge *p=adj[x];p;p=p->nxt) {
                    if( p->to == f || p->to == hvy[x] || p->to == pos[x] ) continue;
                    dp[x] = add(dp[x], dfs3(p->to, x, val[p->to], inv[p->to]));
                }
            }
            else {
                int cnt = 0; bool flag = false;
                for(edge *p=adj[x];p;p=p->nxt) {
                    if( p->to == f ) continue;
                    if( dp[p->to] == 0 ) {
                        cnt++;
                        if( p->to == hvy[x] )
                            flag = true;
                    }
                }
                if( cnt == 2 ) {
                    if( flag ) {
                        T.update2(spro[x]);
                        for(edge *p=adj[x];p;p=p->nxt) {
                            if( p->to == f ) continue;
                            if( dp[p->to] == 0 && p->to != hvy[x] )
                                dp[x] = add(dp[x], dfs3(p->to, x, val[p->to], 1));
                        }
                    }
                    else {
                        T.update2(0);
                        for(edge *p=adj[x];p;p=p->nxt) {
                            if( p->to == f ) continue;
                            if( dp[p->to] == 0 ) {
                                if( !flag ) {
                                    dfs4(p->to, x, val[x] + val[p->to], spro[x]);
                                    flag = true;
                                }
                                else dp[x] = add(dp[x], dfs3(p->to, x, val[p->to], 1));
                            }
                        }
                    }
                }
                T.update2(0);
            }
        }
    }
    dp[x] = add(dp[x], T.query(0)), inv[x] = pow_mod(dp[x], MOD-2);
    if( !flag ) T.init();
}
void solve() {
    int n; scanf("%d", &n);
    for(int i=1;i<=n;i++)
        scanf("%d", &val[i]), adj[i] = NULL;
    ecnt = &edges[0];
    for(int i=1;i<n;i++) {
        int x, y; scanf("%d%d", &x, &y);
        addedge(x, y);
    }
    int rt = T.get_rand() % n + 1;
    dfs1(rt, 0), dfs2(rt, 0, false);
    printf("%d\n", dp[rt]);
}
int main() {
    int t; scanf("%d", &t);
    srand(20041112^t);
    while( t-- ) solve();
}

@details@

好久没写非旋 treap 了（？写过吗），在 merge 的时候写成了按根的权值决定左右儿子关系。。。
因为只有堆才满足根的权值最大，可以只比较根的权值。。。要是两棵平衡树根的权值一样不就炸了。。。
正确写法是在 merge 时不考虑权值，在调用 merge 时保证两棵平衡树的相对大小关系即可。