Submission #1036513

#	Time	Username	Problem	Language	Result	Execution time	Memory
1036513		andrei_c1	Paths (RMI21_paths)	C++17	68 / 100	1045 ms	22356 KiB

Main

#pragma GCC optimize("Ofast")
#pragma GCC target("avx2")
#include <bits/stdc++.h>
using namespace std;
using ll = long long;

const int kN = 1e5;
const int kNil = -1;
const int kBuff = 1 << 16;

int pos, len;
char buf[kBuff];

int nc() {
	if(pos == len) {
		pos = 0;
		len = (int)fread(buf, 1, kBuff, stdin);
		if(!len) {
			return EOF;
		}
	}
	return buf[pos++];
}

int readint() {
	int x;
	char ch;
	while(!isdigit(ch = nc()));
	x = ch - '0';
	while(isdigit(ch = nc())) {
		x = x * 10 + ch - '0';
	}
	return x;
}

int n, k;
vector<pair<int, int>> adj[kN];
vector<int> ord;
int up[kN], val[kN], tin[kN], tout[kN];

void dfs(int u) {
	ord.emplace_back(u);
	tin[u] = ord.size() - 1;
	for(const auto &[to, c]: adj[u]) if(to != up[u]) {
		up[to] = u;
		val[to] = c;
		dfs(to);
	}
	tout[u] = ord.size() - 1;
}

struct node_t {
	ll val;
	int idx;
	node_t() : val(0), idx(0) {}
	node_t(ll val, int idx) : val(val), idx(idx) {}
	node_t operator + (const node_t &oth) const {
		if(val > oth.val) {
			return node_t(val, idx);
		} else {
			return oth;
		}
	};
};

vector<ll> mars;

struct aint {
	int n;
	vector<node_t> tree;
	vector<ll> lazy;
	aint() {}
	aint(int n) : n(n), tree(n << 2 | 1), lazy(n << 2 | 1) {}
	void build(int node, int l, int r) {
		if(l == r) {
			tree[node] = node_t(mars[l], l);
			lazy[node] = 0;
		} else {
			int mid = (l + r) >> 1;
			build(node << 1, l, mid);
			build(node << 1 | 1, mid + 1, r);
			tree[node] = tree[node << 1] + tree[node << 1 | 1];
		}
	}
	void build() {
		build(1, 0, n - 1);
	}
	void push(int node) {
		if(lazy[node]) {
			for(const auto &it: {node << 1, node << 1 | 1}) {
				tree[it].val += lazy[node];
				lazy[it] += lazy[node];
			}
			lazy[node] = 0;
		}
	}
	void update(int a, int b, int val, int node, int l, int r) {
		if(a <= l && r <= b) {
			tree[node].val += val;
			lazy[node] += val;
		} else {
			push(node);
			int mid = (l + r) >> 1;
			if(a <= mid) {
				update(a, b, val, node << 1, l, mid);
			}
			if(b > mid) {
				update(a, b, val, node << 1 | 1, mid + 1, r);
			}
			tree[node] = tree[node << 1] + tree[node << 1 | 1];
		}
	}
	void update(int a, int b, int val) {
		update(a, b, val, 1, 0, n - 1);
	}
	inline node_t all() {
		return tree[1];
	}
	void clear() {
		fill(tree.begin(), tree.end(), node_t());
		fill(lazy.begin(), lazy.end(), 0);
	}
} ds;

ll solve(int root) {
	ll res = 0;
	up[root] = kNil;
	val[root] = 0;
	dfs(root);
	for(int i = 0; i < n; i++) {
		mars[tin[i]] += val[i];
		mars[tout[i] + 1] -= val[i];
	}
	for(int i = 1; i < n; i++) {
		mars[i] += mars[i - 1];
	}
	ds.build();
	vector<bool> vis(n);
	for(int i = 0; i < k; i++) {
		auto [sum, idx] = ds.all();
		if(sum == 0) {
			break;
		}
		res += sum;
		int u = ord[idx];
		while(u != kNil && !vis[u]) {
			vis[u] = true;
			ds.update(tin[u], tout[u], -val[u]);
			u = up[u];
		}
	}
	fill(mars.begin(), mars.end(), 0);
	ds.clear();
	ord.clear();
	return res;
}

ll dist1[kN], dist2[kN];

void dfs(int u, ll dist[], int v = -1) {
	for(const auto &[to, c]: adj[u]) if(to != v) {
		dist[to] = dist[u] + c;
		dfs(to, dist, u);
	}
}

int main() {
	cin.tie(nullptr)->sync_with_stdio(false);
	n = readint();
	k = readint();
	mars = vector<ll>(n + 1);
	ds = aint(n);
	for(int i = 1; i < n; i++) {
		int u, v, c;
		u = readint() - 1;
		v = readint() - 1;
		c = readint();
		adj[u].emplace_back(v, c);
		adj[v].emplace_back(u, c);
	}
	if(k == 1) {
		dist1[0] = 0;
		dfs(0, dist1);
		int x = 0, y = 0;
		for(int i = 1; i < n; i++) {
			if(dist1[x] < dist1[i]) {
				x = i;
			}
		}
		dist2[x] = 0;
		dfs(x, dist2);
		for(int i = 1; i < n; i++) {
			if(dist2[y] < dist2[i]) {
				y = i;
			}
		}
		dist1[y] = 0;
		dfs(y, dist1);
		for(int i = 0; i < n; i++) {
			cout << max(dist1[i], dist2[i]) << '\n';
		}
	} else {
		for(int i = 0; i < n; i++) {
			cout << solve(i) << '\n';
		}
	}
	return 0;
}

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