【HDU5877】Weak Pair(线段树+dfs+离散化)

记录一个菜逼的成长。。

题目大意：
给你一个有根树，每个节点有权值，如果u是v的祖先，并且a[u] * a[v] <= k则（u,v）是一个虚弱对，求有多少组这样的对。

dfs序，dfs到某一节点时，路径上走过的点都是此节点的祖先。

///#pragma comment(linker, "/STACK:1024000000,1024000000")
#include <cstdio>
#include <iostream>
#include <cstring>
#include <string>
#include <algorithm>
#include <cstdlib>
#include <vector>
#include <set>
#include <map>
#include <queue>
#include <stack>
#include <list>
#include <deque>
#include <cctype>
#include <bitset>
#include <cmath>
using namespace std;
#define ALL(v) (v).begin(),(v).end()
#define cl(a,b) memset(a,b,sizeof(a))
#define bp __builtin_popcount
#define pb push_back
#define fin freopen("D://in.txt","r",stdin)
#define fout freopen("D://out.txt","w",stdout)
#define lson t<<1,l,mid
#define rson t<<1|1,mid+1,r
#define seglen (node[t].r-node[t].l+1)
#define pi 3.1415926
#define e  2.718281828459
typedef long long LL;
typedef unsigned long long ULL;
typedef pair<int,int> PII;
typedef pair<LL,LL> PLL;
typedef vector<PII> VPII;
const int INF = 0x3f3f3f3f;
const int MOD = 1e9 + 7;
template <typename T>
inline void read(T &x){
    T ans=0;
    char last=' ',ch=getchar();
    while(ch<'0' || ch>'9')last=ch,ch=getchar();
    while(ch>='0' && ch<='9')ans=ans*10+ch-'0',ch=getchar();
    if(last=='-')ans=-ans;
    x = ans;
}
inline bool DBread(double &num)
{
    char in;double Dec=0.1;
    bool IsN=false,IsD=false;
    in=getchar();
    if(in==EOF) return false;
    while(in!='-'&&in!='.'&&(in<'0'||in>'9'))
        in=getchar();
    if(in=='-'){IsN=true;num=0;}
    else if(in=='.'){IsD=true;num=0;}
    else num=in-'0';
    if(!IsD){
        while(in=getchar(),in>='0'&&in<='9'){
            num*=10;num+=in-'0';}
    }
    if(in!='.'){
        if(IsN) num=-num;
            return true;
    }else{
        while(in=getchar(),in>='0'&&in<='9'){
                num+=Dec*(in-'0');Dec*=0.1;
        }
    }
    if(IsN) num=-num;
    return true;
}
template <typename T>
inline void write(T a) {
    if(a < 0) { putchar('-'); a = -a; }
    if(a >= 10) write(a / 10);
    putchar(a % 10 + '0');
}
/******************head***********************/
const int maxn = 100000 + 10;
int sum[maxn<<3];//数组开小了竟然判TLE  orz..找了好几个小时。。
LL a[maxn<<1],b[maxn<<1],k,ans;
int deep[maxn<<1];
int n,m;
int cnt,head[maxn];
//边数较大，前向星建边
struct Edge{
    int to,next;
}edge[maxn<<1];
void add(int u,int v)
{
    edge[cnt].to = v;
    edge[cnt].next = head[u];
    head[u] = cnt++;
}
//线段树基本操作。
void pushup(int t)
{
    sum[t] = sum[t<<1] + sum[t<<1+1];
}
void build(int t,int l,int r)
{
    if(l == r){
        sum[t] = 0;
        return ;
    }
    int mid = (r + l) >> 1;
    build(lson);
    build(rson);
    pushup(t);
}
int query(int t,int l,int r,int ql,int qr)
{
    if(l >= ql && r <= qr){
        return sum[t];
    }
    int mid = (l + r) >> 1;
    if(ql > mid)return query(t<<1|1,mid + 1,r,ql,qr);
    else if(qr <= mid)return query(t<<1,l,mid,ql,qr);
    else {
        return query(t<<1,l,mid,ql,mid) + query(t<<1|1,mid+1,r,mid+1,qr);
    }
}
void update(int t,int l,int r ,int pos,LL v)
{
    if(l== r){
        sum[t] += v;
        return ;
    }
    int mid = (l + r) >> 1;
    if(pos > mid)update(t<<1|1,mid+1,r,pos,v);
    else if(pos <= mid)update(t<<1,l,mid,pos,v);
    pushup(t);
}
void dfs(int u)
{
    int l = lower_bound(b+1,b+1+m,k/a[u]) - b;
    ans += 1LL * query(1,1,m,1,l);//查询此节点之前有几个满足条件
    int pos = lower_bound(b+1,b+1+m,a[u]) - b;
    update(1,1,m,pos,1);//标记此节点，说明经过此节点
    for( int i = head[u]; ~i; i = edge[i].next)dfs(edge[i].to);
    update(1,1,m,pos,-1);//清除标记
}
int main()
{
    int T;scanf("%d",&T);
    while(T--){
        cl(deep,0);cl(head,-1);
        cl(sum,0);
        ans = cnt = 0;
        scanf("%d%lld",&n,&k);
        for( int i = 1; i <= n; i++ ){
            scanf("%lld",a+i);
            b[i] = a[i];
        }
        m = n;
        //将k / a[i]离散化
        for( int i = 1; i <= n; i++ ){
            b[++m] = k / a[i];
        }
        sort(b+1,b+1+m);
        m = unique(b+1,b+1+m) - (b+1);
        build(1,1,m);
        for( int u,v,i = 0; i < n - 1; i++ ){
            scanf("%d%d",&u,&v);
            add(u,v);
            deep[v]++;
        }
        for( int i = 1; i <= n; i++ ){
            if(deep[i] == 0)dfs(i);//从根节点开始
        }
        printf("%lld\n",ans);
    }
    return 0;
}