【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.5dp[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());
        }
    }
}

【Codeforces 161D】Distance in Tree

猜你喜欢