【Codeforces 161D】Distance in Tree

【链接】 我是链接,点我呀:)

【题意】

问你一棵树上有多少条长度为k的路径

【题解】

树形dp
设
size[i]表示以节点i为根节点的子树的节点个数
dp[i][k]表示以i为根节点的子树里面距离节点i的距离为k的节点有多少个.
长度为k的路径有两种情况.
1.这个路径从x开始,只经过x的一个子树.
2.这个路径经过x,横跨x的两个子树.
第一种情况直接累加dp[x][k]即可
第二种情况,可以枚举其中一棵子树y包括的路径长度i,另外一棵子树(其余size[x]-1棵子树都可以作为另外一棵子树)
自然包括的路径长度就是k-i了
累加0.5*dp[y][i-1]*(dp[x][k-i]-dp[y][k-i-1]);
乘0.5是因为把(x,y)和(y,x)都算一遍
(dp[x][k-i]-dp[y][k-i-1])表示的是除去y子树其余size[x]-1棵子树中和x距离为k-i的节点个数
具体的代码中有注释

【代码】

import java.io.*;
import java.util.*;

public class Main {

	static InputReader in;
	static PrintWriter out;

	public static void main(String[] args) throws IOException{
		//InputStream ins = new FileInputStream("E:\\rush.txt");
		InputStream ins = System.in;
		in = new InputReader(ins);
		out = new PrintWriter(System.out);
		//code start from here
		new Task().solve(in, out);
		out.close();
	}

	static int N = 50000;
	static int K = 500;
	static class Task{
		long [][]dp;
		long ans;
		ArrayList []g;
		int n,k;

		void dfs(int x,int fa) {
			dp[x][0] = 1;
			int len = g[x].size();
			for (int i = 0;i < len;i++) {
				int y = (int)g[x].get(i);
				if (y==fa) continue;
				dfs(y,x);
				for (int j = 1;j <= k;j++) {
					dp[x][j] = dp[x][j] + dp[y][j-1];
				}
			}
			//以x为开始的链
			ans = ans + dp[x][k];

			//跨过两颗子树
			long ans2 = 0;
			if (k>=2)//只有长度大于等于2才可能跨过两棵子树
				for (int i = 0;i < len;i++) {
					int y = (int)g[x].get(i);
					if (y==fa) continue;
					//以子树y的某个节点作为起点,然后到达x再到达其他位置
					for (int l = 1;l<=k-1;l++) {//不能全都在这棵子树里面,得留k-l在另外一个子树
						//l是x到y下面的某个节点的长度,比如l=1那么就是到达y节点
						//从y的话就是往下l-1长度的点
						long temp1 = dp[y][l-1];

						//紧接着要找从x出发长度为k-l的点(不能包括y这棵子树)
						long temp2 = dp[x][k-l]-dp[y][k-l-1];
						ans2 = ans2 + temp1*temp2;
					}
				}
			ans2/=2;//每种都会重复算一次
			ans = ans + ans2;
		}

		public void solve(InputReader in,PrintWriter out) {
			dp = new long[N+10][K+10];
			g = new ArrayList[N+10];
			for (int i = 1;i<= N;i++) g[i] = new ArrayList();

			n = in.nextInt(); k = in.nextInt();

			for (int i = 1;i <= n-1;i++) {
				int x,y;
				x = in.nextInt();y = in.nextInt();
				g[x].add(y);g[y].add(x);
			}

			dfs(1,-1);
			out.println(ans);
		}
	}

	static class InputReader{
		public BufferedReader br;
		public StringTokenizer tokenizer;

		public InputReader(InputStream ins) {
			br = new BufferedReader(new InputStreamReader(ins));
			tokenizer = null;
		}

		public String next(){
			while (tokenizer==null || !tokenizer.hasMoreTokens()) {
				try {
				tokenizer = new StringTokenizer(br.readLine());
				}catch(IOException e) {
					throw new RuntimeException(e);
				}
			}
			return tokenizer.nextToken();
		}

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