HDU - 5877 Weak Pair （DFS+二分+树状数组+离散化）

You are given a rooted tree of N nodes, labeled from 1 to N. To the ith 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≤10^5

  0≤ai≤10^9

  0≤k≤10^18

Output

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

Sample Input

Sample Output

题解：

DFS，每到一个点，就计算出当前点与祖先中满足的对数（树状数组维护）然后把这个点的值（离散化后的）在树状数组中+1

，但是因为点权值会有零，而树状数组无法维护零，所以我们用一个全局变量（zero）来维护祖先中0点的数量。

注意点：

1.根节点需要自己找。

2.DFS从子节点返回后记得删除之前进行的更改。

3.K是long long

代码：

#include <bits/stdc++.h>

using namespace std;

const int MAXN = 100005;

long long N,K,len,zero,ans; 
long long board[MAXN];//存储每个点的权值 
long long bi[MAXN];//用于二分 
int tree[MAXN];//树状数组 
int idx[MAXN];//存每个点离散后的值 
bool Is[MAXN];//判断点是否是根节点 

struct D{
	long long v,id;
	bool operator < (const D& b)const {
		return v < b.v;
	}
}data[MAXN];

struct Edge{
	int next,to;
}E[MAXN*2];

int head[MAXN],tot;

inline void Add(int from,int to){
	E[++tot].next = head[from];
	E[tot].to = to;
	head[from] = tot;
}

int lowbit(int t)
{
    return t&(-t);
}

void add(int index,int value)//单点更新 
{
    for(int i=index ; i<=N ; i+=lowbit(i))
    {
        tree[i] += value;
    }
}

int getSum(int index)//区间查询 
{
    int sum = 0;
    for(int i=index ; i>0 ; i-=lowbit(i))
    {
        sum += tree[i];
    }
    return sum;
}

int MyBi(long long x){//二分查找第一个小于等于x的值的下标。 
	int l=1,r=len;
	int m;
	while(l <= r){
		m = l + (r-l)/2;
		if(bi[m] <= x)l = m + 1;
		else r = m - 1;
	}
	if(l == 1)return -1;
 	return l-1;
}

void DFS(int x){
	if(board[x] == 0){
		ans += getSum(N) + zero;//0和任何数相乘都满足 
		++zero;
	}
	else {
		int t = MyBi(K/board[x]);
		if(t == -1)ans += zero;
		else ans += getSum(t) + zero;
		add(idx[x],1);
	}
	for(int i=head[x] ; i ; i=E[i].next){
		DFS(E[i].to);
	}
	if(board[x] == 0)--zero;
	else add(idx[x],-1);
}

inline void init(){
	memset(tree,0,sizeof tree);
	memset(Is,true,sizeof Is);
	memset(head,0,sizeof head);
	tot = ans = zero = 0;
}


int main(){
	
	int T;
	scanf("%d",&T);
	while(T--){
		init();
		scanf("%lld %lld",&N,&K);
		for(int i=1 ; i<=N ; ++i){
			scanf("%lld",&board[i]);
			bi[i] = board[i];
			data[i].v = board[i];
			data[i].id = i;
		}
		//------------------------------------------------------
		//离散化点权值。 
		sort(data+1,data+1+N);
		long long t = 1;
		idx[data[1].id] = t;
		for(int i=2 ; i<=N ; ++i){
			if(data[i].v == data[i-1].v)idx[data[i].id] = t;
			else idx[data[i].id] = ++t;
		}
		//------------------------------------------------------
		//通过离散去重点权值来得到二分数组 
		sort(bi+1,bi+1+N);
		len = unique(bi+1,bi+1+N) - (bi+1);
		//------------------------------------------------------
		for(int i=1 ; i<=N-1 ; ++i){
			int a,b;
			scanf("%d %d",&a,&b);
			Add(a,b);
			Is[b] = false;
		}
		for(int i=1 ; i<=N ; ++i){
			if(Is[i]){
				DFS(i);
				printf("%d\n",ans);
				break;
			}
		}
	} 
	
	return 0;
}

HDU - 5877 Weak Pair （DFS+二分+树状数组+离散化）

猜你喜欢