161D - Distance in Tree - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
using LL = long long;
int main(){
	ios::sync_with_stdio(false);cin.tie(0);
	int n, k;
	cin >> n >> k;
	vector<vector<int>> e(n + 1);
	for (int i = 0; i < n - 1; i ++ ){
		int u, v;
		cin >> u >> v;
		e[u].push_back(v);
		e[v].push_back(u);
	}
	
	vector<vector<int>> f(n + 1, vector<int>(k + 1));
	int ans = 0;
	function<void(int, int)> dfs = [&](int u, int fa){
		f[u][0] = 1;
		for (auto v : e[u]){
			if (v == fa) continue;
			dfs(v, u);
			for (int i = 0; i < k; i ++ ){
				ans += f[v][i] * f[u][k - i - 1];
			}
			for (int i = 1; i <= k; i ++ ){
				f[u][i] += f[v][i - 1];
			}
		}
	};
	dfs(1, 0);
	cout << ans << "\n";
	return 0;
}