Submission #891265

#	Time	Username	Problem	Language	Result	Execution time	Memory
891265		fanwen	Paths (RMI21_paths)	C++17	100 / 100	214 ms	23984 KiB

Main

#include <bits/stdc++.h>

using namespace std;

#define fi first
#define se second
#define ll long long
#define file(name)                  \
    if(fopen(name".inp", "r"))      \
        freopen(name".inp", "r", stdin), freopen(name".out", "w", stdout); 

const int MAX = 3e5 + 5;

int n, k;
long long dp[MAX];
vector <pair <int, int>> adj[MAX];

namespace sub4 {
	bool check() {
		return n <= 2000;
	}

	void dfs(int u, int p, vector <long long> &res) {
		for (pair <int, int> tmp : adj[u]) if(tmp.fi != p) {
			int v = tmp.fi, w = tmp.se;
			dfs(v, u, res);
			dp[u] = max(dp[u], dp[v] + w);
		}

		bool oke = false;
		for (pair <int, int> tmp : adj[u]) if(tmp.fi != p) {
			int v = tmp.fi, w = tmp.se;
			if(dp[u] != dp[v] + w || oke) res.emplace_back(dp[v] + w);
			else oke = true;
		}
	}

	void solve() {
		for (int i = 1; i <= n; ++i) {
	    	fill(dp + 1, dp + n + 1, 0);
	    	vector <long long> res;
	    	dfs(i, 0, res);
	    	res.push_back(dp[i]);
	    	sort(res.begin(), res.end(), greater <long long>());
	    	long long sum = 0;
	    	for (int j = 0; j < min((int) res.size(), k); ++j) sum += res[j];
	    	cout << sum << '\n';
	    }
	    exit(0);
	}
}

namespace sub_full {
	bool check() {
		return true;
	}

	multiset <long long> values, max_k;
	long long total, ans[MAX];
	pair <long long, long long> dp[MAX];

	void insert(long long val) {
		if(max_k.empty() || val > *max_k.begin()) {
			total += val;
			max_k.insert(val);
		} else {
			values.insert(val);
		}
	}

	void remove(long long val) {
		if(max_k.count(val)) {
			max_k.erase(max_k.find(val));
			total -= val;
		} else {
			if(!values.count(val)) return;
			values.erase(values.find(val));
		}
	}

	void get_ans() {
		while((int) max_k.size() < k && !values.empty()) {
			total += *values.rbegin(); 
			max_k.insert(*values.rbegin());
			values.erase(values.find(*values.rbegin()));
		}

		while((int) max_k.size() > k) {
			total -= *max_k.begin();
			values.insert(*max_k.begin());
			max_k.erase(max_k.begin());
		}

	}

	void dfs(int u, int p) {
		for (pair <int, int> tmp : adj[u]) if(tmp.fi != p) {
			int v = tmp.fi, w = tmp.se;
			dfs(v, u);
			if(dp[u].fi < dp[v].fi + w) {
				dp[u].se = dp[u].fi;
				dp[u].fi = dp[v].fi + w;
			} else dp[u].se = max(dp[u].se, dp[v].fi + w);
		}

		bool used = false;

		for (pair <int, int> tmp : adj[u]) if(tmp.fi != p) {
			int v, w; tie(v, w) = tmp;
			if(dp[u].fi != dp[v].fi + w || used) insert(dp[v].fi + w);
			else used = true;
		}
	}

	void calc(int u, int p) {
		get_ans();
		ans[u] = (k == 1 ? dp[u].fi : total);

		for (pair <int, int> tmp : adj[u]) if(tmp.fi != p) {
			int v, w; tie(v, w) = tmp;

			long long val1 = dp[v].fi;
			remove(val1 + w); insert(val1);
			long long val2 = (dp[u].fi == dp[v].fi + w ? dp[u].se : dp[u].fi);
			remove(val2); insert(val2 + w);

			if(dp[v].fi < val2 + w) {
				dp[v].se = dp[v].fi;
				dp[v].fi = val2 + w;
			} else dp[v].se = max(dp[v].se, val2 + w);

			calc(v, u);

			insert(val1 + w), remove(val1);
			insert(val2), remove(val2 + w);
		}
	}
	void solve() {
		dfs(1, 0); 
		insert(dp[1].fi);
		calc(1, 0);
		for (int i = 1; i <= n; ++i) cout << ans[i] << "\n";
			exit(0);
	}
}

void you_make_it(void) {
    cin >> n >> k;
    for (int i = 1; i < n; ++i) {
    	int u, v, w; cin >> u >> v >> w;
    	adj[u].emplace_back(v, w);
    	adj[v].emplace_back(u, w);
    }
    if(sub_full::check()) sub_full::solve();
    if(sub4::check()) sub4::solve();
}

signed main() {

#ifdef LOCAL
    freopen("TASK.inp", "r", stdin);
    freopen("TASK.out", "w", stdout);
#endif
    file("paths");
    auto start_time = chrono::steady_clock::now();

    cin.tie(0), cout.tie(0) -> sync_with_stdio(0);

    you_make_it();

    auto end_time = chrono::steady_clock::now();

    cerr << "\nExecution time : " << chrono::duration_cast <chrono::milliseconds> (end_time - start_time).count() << "[ms]" << endl;

    return (0 ^ 0);
}

// Dream it. Wish it. Do it.

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:10:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   10 |         freopen(name".inp", "r", stdin), freopen(name".out", "w", stdout);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:164:5: note: in expansion of macro 'file'
  164 |     file("paths");
      |     ^~~~
Main.cpp:10:49: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   10 |         freopen(name".inp", "r", stdin), freopen(name".out", "w", stdout);
      |                                          ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:164:5: note: in expansion of macro 'file'
  164 |     file("paths");
      |     ^~~~

• Subtask 1: 8 / 8
• Subtask 2: 11 / 11
• Subtask 3: 17 / 17
• Subtask 4: 20 / 20
• Subtask 5: 12 / 12
• Subtask 6: 32 / 32