hdu 5877 Weak Pair （Treap)

链接：http://acm.hdu.edu.cn/showproblem.php?pid=5877

题面;

Weak Pair

Time Limit: 4000/2000 MS (Java/Others)    Memory Limit: 262144/262144 K (Java/Others)
Total Submission(s): 5706    Accepted Submission(s): 1617

Problem Description

You are given a rooted tree of N nodes, labeled from 1 to N . To the i th node a non-negative value ai is assigned.An ordered pair of nodes (u,v) is said to be weak if
  (1) u is an ancestor of v (Note: In this problem a node u is not considered an ancestor of itself);
  (2) au×av≤k .

Can you find the number of weak pairs in the tree?

 

Input

There are multiple cases in the data set.
  The first line of input contains an integer T denoting number of test cases.
  For each case, the first line contains two space-separated integers, N and k , respectively.
  The second line contains N space-separated integers, denoting a1 to aN .
  Each of the subsequent lines contains two space-separated integers defining an edge connecting nodes u and v , where node u is the parent of node v .

  Constrains:
  
  1≤N≤105
  
  0≤ai≤109
  
  0≤k≤1018

 

Output

For each test case, print a single integer on a single line denoting the number of weak pairs in the tree.

 

Sample Input

1 2 3 1 2 1 2

 

Sample Output

1

 

Source

2016 ACM/ICPC Asia Regional Dalian Online

 

 

 

思路： treap板子题，注意并没有规定1为根，根要自己找下，还有a[i]可能为0，此时需要特判下

初始化没清空左右儿子，一直RE，真实自闭

实现代码;

#include<bits/stdc++.h>
using namespace std;
#pragma comment(linker, "/STACK:102400000,102400000")
#define ll long long
#define ls t[x].ch[0]
#define rs t[x].ch[1]
const ll M = 2e5 +10;
const ll inf = 1e18+10;
ll rt,sz,ans,a[M],n,k;
struct node{
    ll ch[2],cnt,siz,val,rd;
}t[M];
vector<ll>g[M];
void up(ll x){
    t[x].siz = t[ls].siz + t[rs].siz+t[x].cnt;
}

void rotate(ll &x,ll d){
    ll son = t[x].ch[d];
    t[x].ch[d] = t[son].ch[d^1];
    t[son].ch[d^1] = x; up(x); up(x=son);
}

void ins(ll &x,ll val){
    if(!x){
        x = ++sz;
        t[x].cnt = t[x].siz = 1;
        t[x].val = val,t[x].rd = rand();
        return ;
    }
    t[x].siz ++;
    if(t[x].val == val){
        t[x].cnt++; return ;
    }
    ll d = t[x].val < val; ins(t[x].ch[d],val);
    if(t[x].rd > t[t[x].ch[d]].rd) rotate(x,d);
}

void del(ll &x,ll val){
    if(!x) return ;
    if(t[x].val == val){
        if(t[x].cnt > 1){
            t[x].cnt--,t[x].siz--;return ;
        }
        bool d = t[ls].rd > t[rs].rd;
        if(ls == 0||rs == 0) x = ls+rs;
        else rotate(x,d),del(x,val);
    }
    else t[x].siz--,del(t[x].ch[t[x].val<val],val);
}

ll rk(ll x,ll val){
    if(!x) return 0;
    if(t[x].val == val) return t[ls].siz+t[x].cnt;
    if(t[x].val > val) return rk(ls,val);
    return rk(rs,val)+t[ls].siz+t[x].cnt;
}

void dfs(ll u,ll f){
    ll num = inf;
    if(a[u]!=0) num = k/a[u];
    ans += rk(rt,num);
    ins(rt,a[u]);
    for(ll i = 0;i < g[u].size();i ++){
        ll v = g[u][i];
        if(v == f) continue;
        dfs(v,u);
    }
    del(rt,a[u]);
}

ll d[M];

void init(){
    for(ll i = 1;i < M;i ++){
        t[i].val = 0;d[i] = 0;t[i].cnt=0,t[i].siz = 0;
        t[i].ch[0] = 0; t[i].ch[1] = 0;
    }
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    ll t,x,y;
    cin>>t;
    while(t--){
        rt = 0,ans = 0,sz = 0;
        init();
        cin>>n>>k;
        for(ll i = 1;i <= n;i ++) cin>>a[i];
        for(ll i = 1;i < n;i ++){
            cin>>x>>y;
            g[x].push_back(y);
            g[y].push_back(x);
            d[y]++;
        }
        for(ll i = 1;i <= n;i ++)
        if(d[i]==0) {
            dfs(i,0); break;
        }
        cout<<ans<<endl;
        for(ll i = 1;i <= n ;i ++) g[i].clear();
    }
}