hdu5877 dfs序+树状数组

题意：

有根树，n个点，点的权重w[i] , K 

求这样的<u , v> 总对数，满足：

     u是v的祖先，且w[u]*w[v] <= K  

做法：dfs序 + 树状数组（求区间<=q的个数）

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.math.BigInteger;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {
	public static void main(String[] args) {
		new HDU5877().solve();
	}
}

class HDU5877 {
	class Wall implements Comparable<Wall> {
		long high;
		int id;

		public Wall(long high, int id) {
			this.high = high;
			this.id = id;
		}

		public int compareTo(Wall o) {
			return Long.compare(high, o.high);
		}
	}

	class Ask implements Comparable<Ask>{
		int left;
		int right;
		long high ;

		Ask(int left, int right, long high) {
			this.high = high ;
			this.left = left;
			this.right = right;
		}

		@Override
		public int compareTo(Ask o) {
			return Long.compare(o.high , high) ;
		}
	}

	InputReader in = new InputReader(System.in);
	PrintWriter out = new PrintWriter(System.out);
	final int N = 100008;
	int n;
	Wall[] walls = new Wall[N];
	Ask[] asks = new Ask[N];
	int[] c = new int[N];

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

	void add(int i) {
		for (; i <= n; i += lowbit(i))
			c[i] += 1;
	}

	int sum(int i) {
		int s = 0;
		for (; i > 0; i -= lowbit(i))
			s += c[i];
		return s;
	}

	int sum(int l, int r) {
		return sum(r) - sum(l - 1);
	}
	
	long[] w = new long[N] ;
	int[] inDe = new int[N] ;
	int[] leftId = new int[N] ;
	int[] rightId = new int[N] ;
	int dfsTime ;
	void dfs(int u , int father){
		leftId[u] = ++dfsTime ;
		for(int i = head[u] ; i != -1 ; i = e[i].next){
			int v = e[i].v ;
			if(v == father) continue ;
			dfs(v , u) ;
		}
		rightId[u] = dfsTime ; 
	}
	
	void solve() {
		
		int t = in.nextInt() ;
		while(t-- > 0){
			n = in.nextInt() ;
			long K = in.nextLong() ;
			for(int i = 1 ; i <= n ; i++) w[i] = in.nextLong() ;
			eid = 0 ;
			Arrays.fill(head , 0 , 1+n , -1) ;
			Arrays.fill(inDe , 0 , 1+n , 0) ;
			for(int i = 1 ; i < n ; i++){
				int u = in.nextInt() ;
				int v = in.nextInt() ;
				addEdge(u , v) ;
				inDe[v]++ ;
			}
			
			int root = -1 ;
			for(int i = 1 ; i <= n ; i++){
				if(inDe[i] == 0 ){
					root = i ; 
					break ; 
				}
			}
			
			dfsTime = 0 ; 
			dfs(root, -1) ;
			
			for (int i = 1; i <= n; i++)
				walls[i] = new Wall(w[i] , leftId[i]) ;
			for (int i = 1; i <= n ; i++)
				asks[i] = new Ask(leftId[i]  , rightId[i] , w[i]) ;

			Arrays.sort(walls, 1, 1 + n) ;
			Arrays.sort(asks, 1, 1 + n) ;
			Arrays.fill(c, 0) ;

			long res = 0 ; 
			int pos = 1 ;
			for (int i = 1; i <= n; i++) {
				if(asks[i].high == 0){
					res += asks[i].right - asks[i].left ;
					continue ;
				}
				long limit = K / asks[i].high ;
				while (pos <= n && walls[pos].high <= limit) {
					add(walls[pos].id) ;
					pos++;
				}

				res += sum(asks[i].left+1 , asks[i].right) ;
			}

			out.println(res) ;
		}
		
		out.flush() ;
	}
	
	
	class Edge{
		int v ; 
		int next ;
	    Edge(int v , int next){
	    	this.v = v ; 
	    	this.next = next ;
		}
	}
	Edge[] e = new Edge[N] ;
	int[] head = new int[N] ;
	int eid ;
	void addEdge(int u , int v){
		e[eid] = new Edge(v, head[u]) ;
		head[u] = eid++ ;
	}

}

class InputReader {
	public BufferedReader reader;
	public StringTokenizer tokenizer;

	public InputReader(InputStream stream) {
		reader = new BufferedReader(new InputStreamReader(stream), 32768);
		tokenizer = new StringTokenizer("");
	}

	private void eat(String s) {
		tokenizer = new StringTokenizer(s);
	}

	public String nextLine() {
		try {
			return reader.readLine();
		} catch (Exception e) {
			return null;
		}
	}

	public boolean hasNext() {
		while (!tokenizer.hasMoreTokens()) {
			String s = nextLine();
			if (s == null)
				return false;
			eat(s);
		}
		return true;
	}

	public String next() {
		hasNext();
		return tokenizer.nextToken();
	}

	public int nextInt() {
		return Integer.parseInt(next());
	}

	public long nextLong() {
		return Long.parseLong(next());
	}

	public double nextDouble() {
		return Double.parseDouble(next());
	}

	public BigInteger nextBigInteger() {
		return new BigInteger(next());
	}

}