Submission #830264

#	Time	Username	Problem	Language	Result	Execution time	Memory
830264		fanwen	Race (IOI11_race)	C++17	0 / 100	4 ms	4948 KiB

race

#include <bits/stdc++.h>

using namespace std;

#define MASK(x) (1LL << (x))
#define BIT(x, i) (((x) >> (i)) & 1)
#define ALL(x) (x).begin(), (x).end()
#define REP(i, n) for (int i = 0, _n = n; i < _n; ++i)
#define FOR(i, a, b) for (int i = (a), _b = (b); i <= _b; ++i)
#define FORD(i, a, b) for (int i = (a), _b = (b); i >= _b; --i)
#define FORE(i, a, b) for (int i = (a), _b = (b); i < _b; ++i)
#define debug(...) "[" << #__VA_ARGS__ ": " << (__VA_ARGS__) << "] "
#define file(name) if (fopen(name".inp", "r")) { freopen(name".inp", "r", stdin); freopen(name".out", "w", stdout); }

template <class A, class B> bool minimize(A &a, B b)  { if (a > b) { a = b; return true; } return false; }
template <class A, class B> bool maximize(A &a, B b)  { if (a < b) { a = b; return true; } return false; }

const int MAXN = 2e5 + 5;
const int MAXK = 1e6 + 5;

int n, k, avail[MAXN], sz[MAXN], cnt[MAXK];
vector <pair <int, int>> adj[MAXN];
vector <int> used;
long long ans;

int dfs_size(int u, int par) {
	sz[u] = 1;
	for (auto [v, _] : adj[u]) if(v != par and !avail[v]) {
		sz[u] += dfs_size(v, u);
	}
	return sz[u];
}

int dfs_centroid(int u, int par, int n) {
	for (auto [v, _] : adj[u]) if(v != par and !avail[v]) {
		if(2 * sz[v] >= n) return dfs_centroid(v, u, n);
	}
	return u;
}

void DFS(int u, int par, int depth, int type) {
	if(depth > k) return;
	if(type) ans += cnt[k - depth];
	else {
		cnt[depth]++;
		used.emplace_back(depth);
	}
	for (auto [v, w] : adj[u]) if(v != par and !avail[v]) {
		DFS(v, u, depth + w, type);
	}
}

void solve(int u, int par) {
	int n = dfs_size(u, par);
	avail[u = dfs_centroid(u, par, n)] = true;
	cnt[0] = 1;
	for (auto [v, w] : adj[u]) if(v != par and !avail[v]) {
		DFS(v, u, w, 1);
		DFS(v, u, w, 0);
	}
	for (auto x : used) cnt[x] = 0;
	used.clear();
	for (auto [v, _] : adj[u]) if(v != par and !avail[v]) solve(v, u);
}

int best_path(int a, int b, int H[][2], int L[]) {
	n = a, k = b;
	REP(i, n - 1) {
		adj[H[i][0] + 1].emplace_back(H[i][1] + 1, L[i]);
		adj[H[i][1] + 1].emplace_back(H[i][0] + 1, L[i]);
	}
	solve(1, 0);
	return (ans == 0 ? -1 : ans);
}

#ifdef LOCAL

#include "race.h"
#include <stdio.h>
#include <stdlib.h>

#define MAX_N 500000

int N, K;
int H[MAX_N][2];
int L[MAX_N];
int solution;

inline
void my_assert(int e) {if (!e) abort();}

void read_input()
{
  int i;
  my_assert(2==scanf("%d %d",&N,&K));
  for(i=0; i<N-1; i++)
    my_assert(3==scanf("%d %d %d",&H[i][0],&H[i][1],&L[i]));
  my_assert(1==scanf("%d",&solution));
}

int main()
{
  int ans;
  read_input();
  ans = best_path(N,K,H,L);
  if(ans==solution)
    printf("Correct.\n");
  else
    printf("Incorrect. Returned %d, Expected %d.\n",ans,solution);

  return 0;
}

#endif

// Dream it. Wish it. Do it.

• Subtask 1: 0 / 9
• Subtask 2: 0 / 12
• Subtask 3: 0 / 22
• Subtask 4: 0 / 57