Submission #560594

#	Time	Username	Problem	Language	Result	Execution time	Memory
560594		prarie	Janjetina (COCI21_janjetina)	C++17	35 / 110	400 ms	16912 KiB

Main

#include <bits/stdc++.h>
 
#define fastio ios_base::sync_with_stdio(0), cin.tie(0)
#define endl '\n'
#define fi first
#define se second
#define sz(v) (int)(v).size()
#define all(v) (v).begin(), (v).end()
#define compress(v) sort(all(v)), (v).erase(unique(all((v))), v.end())
#define pb push_back
#define eb emplace_back
 
using namespace std;
 
using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using ld = long double;
 
template <typename T>
inline void Mx(T &x, T y) { x < y && (x = y); }
template <typename T>
inline void Mn(T &x, T y) { x > y && (x = y); }


template <typename T>
struct Fenwick {
	int n;
	vector<T> tree;

	Fenwick(int _n) : n(_n) {
		tree.resize(n);
	}

	void update(int x, T v) {
		for (int i = x; i < n; i += i & -i)
			tree[i] += v;
	}

	T query(int x) {
		T ret = 0;
		for (int i = x; i > 0; i -= i & -i)
			ret += tree[i];
		return ret;
	}

	T query(int l, int r) {
		return query(r) - query(l - 1);
	}
};

const int MX = 1e5 + 5;

int N, K;
ll ans;
vector<pii> adj[MX];

int t_siz[MX];
bool fin[MX];
Fenwick<int> tree_cnt(MX);
int qn;
array<int, 3> que[MX];

int getCent(int cur, int pv, int t_size) {
	for (auto &nxt : adj[cur]) {
		if (nxt.first == pv || fin[nxt.first]) {
			continue;
		}
		if (t_size < 2 * t_siz[nxt.first]) {
			return getCent(nxt.first, cur, t_size);
		}
	}
	return cur;
}
 
void getSz(int cur, int pv) {
	t_siz[cur] = 1;
	for (auto &nxt : adj[cur]) {
		if (nxt.first == pv || fin[nxt.first]) {
			continue;
		}
		getSz(nxt.first, cur);
		t_siz[cur] += t_siz[nxt.first];
	}
}

int dfs(int cur, int pv, int type, int mx, int d) {
	int ret = 1;
	que[qn++] = { mx, d, type };
	tree_cnt.update(d, 1);
	for (auto &nxt : adj[cur]) {
		if (nxt.first == pv || fin[nxt.first]) {
			continue;
		}
		ret += dfs(nxt.first, cur, type, max(mx, nxt.second), d + 1);
	}
	return ret;
}

void solve(int node) {
	getSz(node, -1);
	int cent = getCent(node, -1, t_siz[node]);
	fin[cent] = true;
	qn = 0;
	vector<int> S;
	vector<vector<int>> Z;
	for (auto &nxt : adj[cent]) {
		if (fin[nxt.first]) {
			continue;
		}
		S.pb(dfs(nxt.first, -1, sz(S), nxt.second, 1));
		Z.pb(vector<int>());
		int z = sz(Z) - 1;
		int s = S.back();
		for (int i = qn - s; i < qn; i++) {
			Z[z].pb(que[i][1]);
		}
		sort(all(Z[z]));
	}
	sort(que, que + qn);
	for (int i = qn - 1; i >= 0; i--) {
		auto &[mx, d, type] = que[i];

		int T = mx - (d + K);
		if (T >= 0) {
			ans += 1;
			if (T > 0) {
				ans += tree_cnt.query(1, T);
				ans -= upper_bound(all(Z[type]), T) - Z[type].begin();
			}
		}
		tree_cnt.update(d, -1);
		Z[type].pop_back();
	}
	for (auto &nxt : adj[cent]) {
		if (fin[nxt.first]) {
			continue;
		}
		solve(nxt.first);
	}
}

int main() {
	fastio;
	cin >> N >> K;
	for (int i = 0; i < N - 1; i++) {
		int u, v, w;
		cin >> u >> v >> w;
		adj[u].eb(v, w);
		adj[v].eb(u, w);
	}

	solve(1);
	cout << ans * 2 << endl;
	return 0;
}

• Subtask 1: 0 / 15
• Subtask 2: 35 / 35
• Subtask 3: 0 / 60