Submission #1159781

#	Time	Username	Problem	Language	Result	Execution time	Memory
1159781		hmm789	Janjetina (COCI21_janjetina)	C++20	0 / 110	2 ms	5264 KiB

Main

#include <bits/stdc++.h>
using namespace std;
#define int long long 

int fw[100005], k;
vector<int> loc;
void update(int x, int v) {
	for(x++; x < 100005; x += x&-x) {fw[x] += v; loc.push_back(x);}
}
int query(int x) {
	int ans = 0;
	for(x++; x; x -= x&-x) ans += fw[x];
	return ans;
}
int query(int x, int y) {
	return query(y) - query(x-1);
}
void reset() {
	for(int i : loc) fw[i] = 0;
	loc.clear();
}

vector<pair<int, int>> adj[100000];
int sz[100000], used[100000], ans;
vector<pair<int, int>> v[100000];

int size(int x, int pa) {
	sz[x] = 1;
	for(auto i : adj[x]) if(i.first != pa && !used[i.first]) sz[x] += size(i.first, x);
	return sz[x];
}
int centroid(int x, int pa, int n) {
	for(auto i : adj[x]) if(i.first != pa && !used[i.first]) {
		if(sz[i.first] > n/2) return centroid(i.first, x, n);
	}
	return x;
}

void dfs(int x, int p, int d, int mx, int pp) {
	v[pp].push_back({mx, d});
	for(auto i : adj[x]) if(i.first != p && !used[i.first]) dfs(i.first, x, d+1, max(mx, i.second), pp); 
}

void decomp(int x) {
	int c = centroid(x, -1, size(x, -1));
	for(auto i : adj[c]) if(!used[i.first]) dfs(i.first, c, 1, i.second, i.first);
	vector<pair<int, int>> v2;
	v2.push_back({0, 0});
	for(auto i : adj[c]) if(!used[i.first]) {
		sort(v[i.first].begin(), v[i.first].end());
		for(auto j : v[i.first]) {
			if(j.first-j.second-k >= 0) ans -= query(j.first-j.second-k);
			update(j.second, 1);
			v2.push_back(j);
		}
		reset();
	}
	sort(v2.begin(), v2.end());
	for(auto j : v2) {
		if(j.first-j.second-k >= 0) ans += query(j.first-j.second-k);
		update(j.second, 1);
	}
	reset();
	used[c] = 1;
	for(auto i : adj[c]) if(!used[i.first]) decomp(i.first);
}

int32_t main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0); cout.tie(0);
	int n, x, y, w;
	cin >> n >> k;
	for(int i = 1; i < n; i++) {
		cin >> x >> y >> w;
		x--; y--;
		adj[x].push_back({y, w});
		adj[y].push_back({x, w});
	}
	decomp(0);
	cout << ans*2 << '\n';
}

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