Submission #446004

#	Time	Username	Problem	Language	Result	Execution time	Memory
446004		grt	Janjetina (COCI21_janjetina)	C++17	110 / 110	979 ms	23552 KiB

Main

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/trie_policy.hpp>
#define ST first
#define ND second
#define PB push_back

using namespace std;
using namespace __gnu_pbds;
using ll = long long;
using pi = pair<int,int>;
using vi = vector<int>;

typedef tree<pi,null_type,less<pi>,rb_tree_tag,tree_order_statistics_node_update>ordered_set;

const int nax = 100 * 1000 + 10, INF = 1e9;
int n, k;
vector<pi>V[nax];
bool mark[nax];
int dp[nax];
vi all;
ll ans;
ordered_set S;

void dfs(int x, int p) {
	dp[x] = 1;
	all.PB(x);
	for(auto nbh : V[x]) if(nbh.ST != p && !mark[nbh.ST]) {
		dfs(nbh.ST, x);
		dp[x] += dp[nbh.ST];
	}
}

vector<pi>path;

void dfsC(int x, int p, int l, int mx) {
	path.emplace_back(mx, l);
	for(auto nbh : V[x]) if(nbh.ST != p && !mark[nbh.ST]) {
		dfsC(nbh.ST, x, l + 1, max(mx, nbh.ND));
	}
}


void centro(int x) {
	all.clear();
	path.clear();
	dfs(x, 0);
	int c = x;
	for(int y : all) if(dp[y] < dp[c] && 2 * dp[y] >= dp[x]) c = y;
	dfsC(c, 0, 0, 0);
	sort(path.begin(), path.end());
	int id = 1;
	for(auto [w, l] : path) {
		//cerr << c << ": " << w << " " << l << " " << S.order_of_key({w - l - k + 1, INF}) << "\n";
		ans += S.order_of_key({w - l - k, INF});
		S.insert({l, id++});
	}
	//cout << ans << "\n";
	S.clear();
	mark[c] = true;
	for(auto nbh : V[c]) if(!mark[nbh.ST]) {
		id = 1;
		path.clear();
		dfsC(nbh.ST, c, 1, nbh.ND);
		sort(path.begin(), path.end());
		for(auto [w, l] : path) {
			ans -= S.order_of_key({w - l - k, INF});
			S.insert({l,id++});
		}
		S.clear();
		centro(nbh.ST);
	}
}

int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	cin >> n >> k;
	for(int a, b, c, i = 1; i < n; ++i) {
		cin >> a >> b >> c;
		V[a].emplace_back(b, c);
		V[b].emplace_back(a, c);
	}
	centro(1);
	cout << ans * 2;
}

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