【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());
		}
	}
}