Codeforces 486D Valid Sets 樹形dp+計數

Codeforces 486D Valid Sets 樹形dp+計數

題目鏈接：點擊打開鏈接

題意：給定常數d 和 n，表示n個點的樹，每個點有點權

問：有多少個點集使得點集中的點兩兩可達且最大點權-最小點權<=d

思路:

對於每個點，計算包含這個點且這個點作為點集中的最小點權時的個數，枚舉每個點。

import java.io.PrintWriter;
import java.text.DecimalFormat;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.Deque;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Scanner;
import java.util.Stack;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.Queue;

public class Main {	
	int n, d;
	int[] a = new int[N];
	int ad(int x, int y){
		x += y;
		if(x>=mod)x %= mod;
		return x;
	}
	int mul(int x, int y){
		long xx = (long)x*(long)y;
		if(xx>=mod) xx %= mod;
		return (int)xx;
	}
	int dfs(int u, int fa, int dep){
		int ans = 1;
		for(int i = head[u]; i!=-1; i = edge[i].nex){
			int v = edge[i].to;
			if((v == fa)||(a[v]>a[dep]||(a[v]==a[dep]&&v>dep))||(a[dep]-a[v]>d))
				continue;			
			ans = mul(ans, dfs(v, u, dep)+1);
		}
		return ans;
	}
	
	void input(){
		d = cin.nextInt(); n = cin.nextInt();
		for(int i = 1; i <= n; i++)	a[i] = cin.nextInt();		
		init_edge();
		for(int i = 1, u, v; i < n; i++){
			u = cin.nextInt(); v = cin.nextInt();
			add(u, v); add(v, u);
		}
	}
	
	void work()	{
		input();
		int ans = 0;
		for(int i = 1; i <= n; i++)
			ans = ad(ans, dfs(i, i, i));
		out.println(ans%mod);
	}
	
	Main() {
		cin = new Scanner(System.in);
		out = new PrintWriter(System.out);
	}

	public static void main(String[] args) {
		Main e = new Main();
		e.work();
		out.close(); 
	}

	public Scanner cin;
	public static PrintWriter out;
	static int N = 2005;
	static int M = 2005;
	DecimalFormat df=new DecimalFormat("0.0000");
	static int inf = (int) 1e9 + 7;
	static long inf64 = (long) 1e18;
	static double eps = 1e-8;
	static double Pi = Math.PI;
	static int mod = 1000000007 ;

	class Edge{
		int from, to, nex;
		Edge(){}
		Edge(int from, int to, int nex){
			this.from = from;
			this.to = to;
			this.nex = nex;
		}
	}
	Edge[] edge = new Edge[M<<1];
	int[] head = new int[N];
	int edgenum;
	void init_edge(){for(int i = 0; i < N; i++)head[i] = -1; edgenum = 0;}
	void add(int u, int v){
		edge[edgenum] = new Edge(u, v, head[u]);
		head[u] = edgenum++;
	}

	int Pow(int x, int y) {
		int ans = 1;
		while (y > 0) {
			if ((y & 1) > 0)
				ans *= x;
			y >>= 1;
			x = x * x;
		}
		return ans;
	}
	double Pow(double x, int y) {
		double ans = 1;
		while (y > 0) {
			if ((y & 1) > 0)
				ans *= x;
			y >>= 1;
			x = x * x;
		}
		return ans;
	}
	int Pow_Mod(int x, int y, int mod) {
		int ans = 1;
		while (y > 0) {
			if ((y & 1) > 0)
				ans *= x;
			ans %= mod;
			y >>= 1;
			x = x * x;
			x %= mod;
		}
		return ans;
	}
	long Pow(long x, long y) {
		long ans = 1;
		while (y > 0) {
			if ((y & 1) > 0)
				ans *= x;
			y >>= 1;
			x = x * x;
		}
		return ans;
	}
	long Pow_Mod(long x, long y, long mod) {
		long ans = 1;
		while (y > 0) {
			if ((y & 1) > 0)
				ans *= x;
			ans %= mod;
			y >>= 1;
			x = x * x;
			x %= mod;
		}
		return ans;
	}

	int gcd(int x, int y){
		if(x>y){int tmp = x; x = y; y = tmp;}
		while(x>0){
			y %= x;
			int tmp = x; x = y; y = tmp;
		}
		return y;
	}
	int max(int x, int y) {
		return x > y ? x : y;
	}

	int min(int x, int y) {
		return x < y ? x : y;
	}

	double max(double x, double y) {
		return x > y ? x : y;
	}

	double min(double x, double y) {
		return x < y ? x : y;
	}

	long max(long x, long y) {
		return x > y ? x : y;
	}

	long min(long x, long y) {
		return x < y ? x : y;
	}

	int abs(int x) {
		return x > 0 ? x : -x;
	}

	double abs(double x) {
		return x > 0 ? x : -x;
	}

	long abs(long x) {
		return x > 0 ? x : -x;
	}

	boolean zero(double x) {
		return abs(x) < eps;
	}
	double sin(double x){return Math.sin(x);}
	double cos(double x){return Math.cos(x);}
	double tan(double x){return Math.tan(x);}
	double sqrt(double x){return Math.sqrt(x);}
}