Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <bits/stdc++.h>
- int main () {
- int n, maxK;
- std::cin >> n >> maxK;
- std::vector<std::vector<int>> edges(n);
- for (int i = 0; i < n - 1; ++i) {
- int u, v;
- std::cin >> u >> v;
- --u, --v;
- edges[u].push_back(v);
- edges[v].push_back(u);
- }
- std::vector<std::vector<int64_t>> dp(n, std::vector<int64_t>(maxK + 1, 0));
- auto dfs = [&](int cur, int dad, auto&& dfs) -> void {
- for (const auto nei : edges[cur])
- if (nei != dad)
- dfs(nei, cur, dfs);
- auto& dp1 = dp[cur];
- std::vector<int64_t> dp2(maxK + 1, 0);
- dp1[0] = 1;
- for (const auto nei : edges[cur]) {
- if (nei == dad)
- continue;
- std::fill(dp2.begin(), dp2.end(), 0);
- for (int k1 = maxK; k1 >= 0; --k1) {
- if (!dp1[k1])
- continue;
- //we skip this nei -> so we get +1 blue
- if (k1 != maxK)
- dp2[k1 + 1] += dp1[k1];
- for (int k2 = 0; k1 + k2 <= maxK; ++k2)
- dp2[k1 + k2] += dp1[k1] * dp[nei][k2];
- }
- std::swap(dp1, dp2);
- }
- };
- dfs(0, -1, dfs);
- int64_t ans = 1;
- for (int cur = 0; cur < n; ++cur) {
- ans += std::accumulate(dp[cur].begin(), dp[cur].end(), 0LL);
- if (cur)
- ans -= dp[cur][maxK];
- }
- std::cout << ans << "\n";
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment