Submission #860351

#	Time	Username	Problem	Language	Result	Execution time	Memory
860351		serifefedartar	Janjetina (COCI21_janjetina)	C++17	110 / 110	291 ms	16020 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
#define fast ios::sync_with_stdio(0);cin.tie(0);
#define s second
#define f first
typedef long long ll;
const ll MOD = 1e9 + 7;
const ll LOGN = 18; 
const ll MAXN = 1e5 + 100;

vector<vector<pair<int,int>>> graph;
int marked[MAXN], sz[MAXN], fen[MAXN], k;
ll ans = 0;
ll get(int k) {
	k++;
	if (k <= 0)
		return 0;
	ll res = 0;
	while (k) {
		res += fen[k];
		k -= k & -k;
	}
	return res;
}

void update(int k, int v) {
	k++;
	while (k < MAXN) {
		fen[k] += v;
		k += k & -k;
	}
}

int get_sz(int node, int parent) {
	sz[node] = 1;
	for (auto u : graph[node]) {
		if (u.f == parent || marked[u.f])
			continue;
		sz[node] += get_sz(u.f, node);
	}
	return sz[node];
}

int find_centro(int node, int parent, int n) {
	for (auto u : graph[node]) {
		if (u.f != parent && !marked[u.f] && sz[u.f] * 2 >= n)
			return find_centro(u.f, node, n);
	}
	return node;
}

vector<pair<int,int>> v;
void dfs(int node, int parent, int mx, int len) {
	v.push_back({mx, len});
	for (auto u : graph[node]) {
		if (u.f != parent && !marked[u.f])
			dfs(u.f, node, max(mx, u.s), len+1);
	}
}

ll calc(int node, int parent, int mx, int len) {
	v.clear();
	dfs(node, parent, mx, len);
	sort(v.begin(), v.end());

	ll ret = 0;
	for (auto u : v) {
		int val = u.f - u.s;
		ret += get(val - k);
		update(u.s, 1);
	}
	for (auto u : v)
		update(u.s, -1); 
	return ret;
}

void decompose(int node, int parent) {
	int n = get_sz(node, parent);
	int centro = find_centro(node, parent, n);

	ans += calc(centro, centro, 0, 0);
	marked[centro] = 1;
	for (auto u : graph[centro]) {
		if (!marked[u.f])
			ans -= calc(u.f, u.f, u.s, 1);
	}

	for (auto u : graph[centro]) {
		if (!marked[u.f])
			decompose(u.f, node);
	}
}

int main() {
	fast
	int n, x, y, w;
	cin >> n >> k;

	graph = vector<vector<pair<int,int>>>(n+1, vector<pair<int,int>>());
	for (int i = 1; i < n; i++) {
		cin >> x >> y >> w;
		graph[x].push_back({y, w});
		graph[y].push_back({x, w});
	}
	decompose(1, 1);
	cout << ans*2 << "\n";
}

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