Submission #1119872

#	Time	Username	Problem	Language	Result	Execution time	Memory
1119872		LonlyR	Janjetina (COCI21_janjetina)	C++17	110 / 110	303 ms	19240 KiB

Main

#include<bits/stdc++.h>
#define int long long
#define ii pair<int,int>

using namespace std;
const int maxn = 1e5 + 5;
int n, k, ans;
int f[maxn], sz[maxn], mark[maxn];
vector<ii> adj[maxn], a;

void upd(int x, int val)
{
    x++;
	for (int i = x; i <= n + 1; i += i & -i)
		f[i] += val;
}

int get(int x)
{
    x++;
	int ans = 0;
	for (int i = x; i > 0; i -= i & -i)
		ans += f[i];
	return ans;
}

void get_size(int x, int p)
{
	sz[x] = 1;
	for (auto [i, j] : adj[x])
		if (i != p && !mark[i])
			get_size(i, x),
			sz[x] += sz[i];
}

int find(int x, int p, int half)
{
	for (auto [i, j] : adj[x])
		if (i != p && !mark[i] && sz[i] > half)
			return find(i, x, half);
	return x;
}

int cal()
{
	sort(a.begin(), a.end());
	int ans = 0;
	for (auto &[mx, len] : a)
	{
		int q = min(n, mx - len - k);
		ans += get(q);
		upd(len, 1);
	}
	for (auto &[mx, len] : a)
		upd(len, -1);
	a.clear();
	return ans;
}

void dfs(int x, int p, int mx, int len)
{
	a.emplace_back(mx, len);
	for (auto [i, j] : adj[x])
		if (i != p && !mark[i])
			dfs(i, x, max(mx, j), len + 1);
}

void solve(int x = 1)
{
	get_size(x, x);
	x = find(x, x, sz[x] / 2);
	mark[x] = 1;
	dfs(x, x, 0, 0);
	ans += cal();
	for (auto [i, j] : adj[x]) if (!mark[i])
		dfs(i, x, j, 1),
		ans -= cal();
	for (auto [i, j] : adj[x])
		if (!mark[i])
			solve(i);
}

signed main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);
//    freopen("test.inp", "r", stdin);
//    freopen("test.out", "w", stdout);
    cin >> n >> k;
    for (int i = 1, u, v, w; i < n; i++)
		cin >> u >> v >> w,
		adj[u].emplace_back(v, w),
		adj[v].emplace_back(u, w);
	solve();
	cout << ans * 2;
}

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