Submission #789837

#	Time	Username	Problem	Language	Result	Execution time	Memory
789837		AmirAli_H1	Janjetina (COCI21_janjetina)	C++17	110 / 110	321 ms	17784 KiB

Main

// In the name of Allah

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

typedef long long int	ll;
typedef long double	ld;
typedef pair<int, int>	pii;
typedef pair<ll, ll>	pll;

#define all(x)		(x).begin(),(x).end()
#define len(x)		((ll) (x).size())
#define F		first
#define S		second
#define pb		push_back
#define sep             ' '
#define endl            '\n'
#define Mp		make_pair
#define debug(x)	cerr << #x << ": " <<  x << endl;
#define kill(x)		cout << x << '\n', exit(0)
#define set_dec(x)	cout << fixed << setprecision(x);
#define file_io(x,y)	freopen(x, "r", stdin); freopen(y, "w", stdout);

int n, k;
const int maxn = 1e5 + 4;
vector<pii> adj[maxn];
bool mark[maxn]; int sz[maxn], val[maxn];
vector<pair<int, pii>> res; ll output = 0;
ll t[maxn];

void add(int i, ll x) {
	for (++i; i < maxn; i += (i & -i)) t[i] += x;
}

ll get(int i) {
	ll res = 0;
	for (++i; i > 0; i -= (i & -i)) res += t[i];
	return res;
}

void dfs(int v, int p = -1) {
	sz[v] = 1;
	for (auto f : adj[v]) {
		int u = f.F, w = f.S;
		if (u == p || mark[u]) continue;
		dfs(u, v);
		sz[v] += sz[u];
	}
}

int get_cent(int v, int n, int p = -1) {
	for (auto f : adj[v]) {
		int u = f.F, w = f.S;
		if (u == p || mark[u]) continue;
		
		if (2 * sz[u] > n) return get_cent(u, n, v);
	}
	return v;
}

void dfs_res(int v, int p = -1, int val = 0, int d = 0) {
	res.pb(Mp(val, Mp(d, v)));
	for (auto f : adj[v]) {
		int u = f.F, w = f.S;
		if (u == p || mark[u]) continue;
		
		dfs_res(u, v, max(val, w), d + 1);
	}
}

ll get_res(int v, int val, int d) {
	res.clear(); dfs_res(v, -1, val, d);
	sort(all(res));
	
	ll ans = 0;
	for (auto f : res) {
		int val = f.F, d = f.S.F, v = f.S.S;
		ans += get(val - d - k); add(d, 1);
	}
	for (auto f : res) {
		int val = f.F, d = f.S.F, v = f.S.S;
		add(d, -1);
	}
	return ans;
}

void solve(int v) {
	dfs(v); v = get_cent(v, sz[v]);
	
	output += get_res(v, 0, 0);
	mark[v] = 1;
	for (auto f : adj[v]) {
		int u = f.F, w = f.S;
		if (mark[u]) continue;
		
		output -= get_res(u, w, 1);
	}
	
	for (auto f : adj[v]) {
		int u = f.F, w = f.S;
		if (mark[u]) continue;
		
		solve(u);
	}
}

int main() {
	ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
	
	cin >> n >> k;
	for (int i = 0; i < n - 1; i++) {
		int u, v, w;
		cin >> u >> v >> w; u--; v--;
		adj[u].pb(Mp(v, w)); adj[v].pb(Mp(u, w));
	}
	
	solve(0);
	
	cout << 2ll * output << endl;
	
	return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void dfs(int, int)':
Main.cpp:44:16: warning: unused variable 'w' [-Wunused-variable]
   44 |   int u = f.F, w = f.S;
      |                ^
Main.cpp: In function 'int get_cent(int, int, int)':
Main.cpp:53:16: warning: unused variable 'w' [-Wunused-variable]
   53 |   int u = f.F, w = f.S;
      |                ^
Main.cpp: In function 'll get_res(int, int, int)':
Main.cpp:77:29: warning: unused variable 'v' [-Wunused-variable]
   77 |   int val = f.F, d = f.S.F, v = f.S.S;
      |                             ^
Main.cpp:81:7: warning: unused variable 'val' [-Wunused-variable]
   81 |   int val = f.F, d = f.S.F, v = f.S.S;
      |       ^~~
Main.cpp:81:29: warning: unused variable 'v' [-Wunused-variable]
   81 |   int val = f.F, d = f.S.F, v = f.S.S;
      |                             ^
Main.cpp: In function 'void solve(int)':
Main.cpp:100:16: warning: unused variable 'w' [-Wunused-variable]
  100 |   int u = f.F, w = f.S;
      |                ^

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