Balkan OI 2017 - City Attractions

Introduction

Let the node that Gigel goes to from node $i$ be $t_i$. Since the graph made from the directed edges $i\to t_i$ is a functional graph, we can use binary jumping or any other efficient method to find the final node.

Now we only need to find all $t_i$ and we are done! However, this isn't as straightforward as it sounds...

A simpler problem

Let's consider a simpler problem: suppose we root the tree at node $1$ and Gigel can only move down the tree (don't worry about the leaves). In this problem, we can find all $t_i$ (and $a[t_i]-dist(i,t_i)$) using a simple tree DP:

Let $dp[i]$ be the node in the subtree of $i$ (excluding $i$ itself) such that $a[dp[i]]-dist(i,dp[i])$ is maximized. Additionally, we store this value in the DP array. We can find $dp[i]$ by taking the best of $c$ and $dp[c]$ over all children $c$ of $i$.

This algorithm runs in $\mathcal{O}(N)$ time.

Finding all $t_i$

Obviously, the solution to the simpler problem doesn't solve the general problem: we might need to move up into a node's parent!

To address this issue, we can first do a DFS to find $dp$ as defined above, and then do a second DFS to allow moving to a node's parent. See the solving for all roots module if you're unfamiliar with this technique. Essentially, we find the best destination from $i$ if we move up into $i$'s parent and then compare that with $dp[i]$.

After doing this, $dp[i]$ is simply $t_i$ as desired.

These two DFSes run in $\mathcal{O}(N)$ time, so the final complexity is $\mathcal{O}(N+\log K)$ (from the binary jumping).

Code

1#include <bits/stdc++.h>
2typedef long long ll;
3using namespace std;
4
5int a[300001], nxt[2][300001];
6pair<int, int> dp[300001], pdp[300001];
7vector<int> graph[300001];
8
9void dfs1(int node = 1, int parent = 0) {
10    dp[node] = {INT_MAX / 2, 0};

Alternative Code (Ben)

1const int MX = 3e5+5;
2
3int N; ll K;
4array<pi,2> bes[MX]; // for each x,
5// get best two a[y]-d(x,y), possibly including x itself
6vi adj[MX];
7int nex[MX], vis[MX];
8
9array<pi,2> comb(array<pi,2> a, array<pi,2> b) {
10    trav(t,b) t.f --;