fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. const int maxn = 200000;
  5. int dp[maxn][2];
  6. vector< pair<int,int> > E[maxn];
  7. int n,k;
  8.  
  9. bool cmp(const pair<int,int> &a, const pair<int,int> &b){
  10. 	return (a.first > b.first);
  11. }
  12.  
  13. void dfs(int u, int p){
  14. 	vector< pair<int,int> > c;
  15. 	set<int> st;
  16.  
  17. 	for(auto e: E[u]){
  18. 		int v = e.first;
  19. 		int cst = e.second;
  20. 		if(v == p) continue;
  21. 		dfs(v,u);
  22. 		c.push_back({dp[v][0] + cst,v});
  23. 	}
  24.  
  25. 	sort(c.begin(), c.end(), cmp);
  26. 	int tk = min(k-1, (int)(c.size()));
  27. 	for(int i = 0; i < tk; i++){
  28. 		dp[u][0] += c[i].first;
  29. 		st.insert(c[i].second);
  30. 	}
  31.  
  32. 	int extra = 0;
  33. 	if(tk != (int)(c.size()))
  34. 		extra = c[tk].first;
  35.  
  36. 	for(auto e: E[u]){
  37. 		int v = e.first;
  38. 		int cst = e.second;
  39. 		if(v == p) continue;
  40.  
  41. 		if(st.count(v) == 0)
  42. 			dp[u][1] = max(dp[u][1], dp[v][1] + dp[u][0] + cst);
  43. 		else
  44. 			dp[u][1] = max(dp[u][1], dp[u][0] - dp[v][0] + extra + dp[v][1]);
  45. 	}
  46. 	st.clear();
  47. 	c.clear();
  48. }
  49.  
  50. int main(){
  51. 	scanf("%d%d", &n, &k);
  52. 	for(int i = 1; i < n; i++){
  53. 		int u,v,c;
  54. 		scanf("%d%d%d", &u, &v, &c);
  55. 		E[u].push_back({v,c});
  56. 		E[v].push_back({u,c});
  57. 	}
  58. 	dfs(0,-1);
  59. 	cout << max(dp[0][0], dp[0][1]) << "\n";
  60. 	return 0;
  61. }
Runtime error #stdin #stdout 0s 22312KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/dYOdvx
language:
C++14 (gcc 8.3)
created:
6 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!