# bvd's Problem of the Week (May 23-May 29)

Problem link: https://codeforces.com/problemset/problem/516/D

## Problem summary

You are given a **weighted** tree with $n$ nodes ($1 \leq n \leq 10^5$).

A leaf is defined as a node with degree 1.

Let $d_i$ be the longest distance from node $i$ to a leaf on a tree.

You are given $q$ queries ($1 \leq q \leq 50$). Each query gives a number $l$ and ask you to return the size of largest set of nodes $S$ such that $\forall i,j \in S: |d_i - d_j| \leq l$ and the set of nodes is connected.

A $O(qn\log^2{n})$ algorithm is needed.

## Part 1: Calculate $d_i$

First, we need to find a fast way to calculate $d_i$. We do this by:
1. Find the diameter $s-t$ of the tree.
2. Find the distance from $s$ to each node on the tree. Find the distance from $t$ to each node on the tree. For each node $i$, the maximum of these two distances is $d_i$

This algorithm works in $O(n)$

Why this works? We have the following theorem:

Given a undirected tree $T = (V,E,w)$ with $w > 0$. $V$ is a set of vertices ($|V| = n$), $E \subset V \times V$ is the set of edges, and $w$ is the weight function ($w: E \rightarrow \mathbb{R}^{+}$)

Let $s, t$ be two endpoints of a diameter of $T$

Prove that a longest path on $T$ from a node $u$ is the path from $u$ to $s$ or the path from $u$ to $t$

<details>
<summary>Proof</summary>

Choose any $u \in V$

Let $u$ be the root of $T$

Assume, for a contradiction, that a longest path from $u$ is neither the path from $u$ to $s$ nor the path from $u$ to $t$.

Then, a longest path from $u$ on $T$ is the path from $u$ to $v$ for some $v \in V, v \ne s, v \ne t$

Let $a = \text{lca}(v,s), b = \text{lca}(s,t)$

$a = \text{lca}(v,s) \Rightarrow \begin{cases} a \rightarrow s \\ a \rightarrow v \end{cases}$

$b = \text{lca}(s,t) \Rightarrow \begin{cases} b \rightarrow s \\ b \rightarrow t \end{cases}$

Because $a \rightarrow s$ and $b \rightarrow s$, either $a \rightarrow b$ or $b \rightarrow a$

Case 1: $a \rightarrow b$

$\begin{cases}
a \rightarrow b \\
b \rightarrow t
\end{cases} \Rightarrow a \rightarrow t$

Because of the assumption, $f(u,v) > f(u,t)$

$u$ is the root of $T \Rightarrow u \rightarrow a$
$\begin{cases}
u \rightarrow a \\
a \rightarrow v
\end{cases} \Rightarrow f(u,a) + f(a,v) = f(u,v)$

$\begin{cases}
u \rightarrow a \\
a \rightarrow t
\end{cases} \Rightarrow f(u,a) + f(a,t) = f(u,t)$

$\begin{cases}
f(u,v) > f(u,t) \\
f(u,a) + f(a,v) = f(u,v) \\
f(u,a) + f(a,t) = f(u,t)
\end{cases} \Rightarrow f(u,a) + f(a,v) > f(u,a) + f(a,t) \Rightarrow f(a,v) > f(a,t)$

$\begin{cases}
a \rightarrow b \\
b \rightarrow t
\end{cases} \Rightarrow f(a,b) + f(b,t) = f(a,t) \Rightarrow f(b,t) \leq f(a,t)$

$\begin{cases}
a \rightarrow b \\
b \rightarrow s
\end{cases} \Rightarrow f(a,b) + f(b,s) = f(a,s) \Rightarrow f(b,s) \leq f(a,s)$

$b = \text{lca}(s,t) \Rightarrow f(s,t) = f(s,b) + f(b,t)$

$a = \text{lca}(v,s) \Rightarrow f(v,s) = f(a,v) + f(a,s)$

We have:

$f(s,t) = f(s,b) + f(b,t) \leq f(a,s) + f(a,t) < f(a,s) + f(a,v) = f(v,s)$

$\Rightarrow f(s,t) < f(v,s)$

$\Rightarrow$ the path from $s$ to $t$ can't be a diameter of $T$, which is absurb

Case 2: $b \rightarrow a$

First, I want to show $b = \text{lca}(v,t)$

$\begin{cases}
b \rightarrow a \\
a \rightarrow v
\end{cases} \Rightarrow b \rightarrow v$

I now have $\begin{cases} b \rightarrow v \\ b \rightarrow t \end{cases}$ To show that $b = \text{lca}(v,t)$, I need to show that $\forall c \in V, c \rightarrow v, c \rightarrow t: c \rightarrow b$ using contradiction.

Assume, for a contradiction, $\exists c \in V, c \rightarrow v, c \rightarrow y: c \nrightarrow b$.

$\begin{cases} c \rightarrow v \\ b \rightarrow v \\ c \nrightarrow b \end{cases} \Rightarrow b \rightarrow c$ and $b \ne c$

$\begin{cases}
a \rightarrow v \\
c \rightarrow v
\end{cases} \Rightarrow$ either $a \rightarrow c$ or $c \rightarrow a$

Suppose that $c \rightarrow a$.

$\begin{cases}
c \rightarrow a \\
a \rightarrow s
\end{cases} \Rightarrow c \rightarrow s$

$\begin{cases}
c \rightarrow s \\
c \rightarrow t \\
b \rightarrow c \\
b \ne c
\end{cases} \Rightarrow b \ne \text{lca}(s,t)$, which is absurb

Suppose that $a \rightarrow c$

$\begin{cases}
a \rightarrow c \\
c \rightarrow t
\end{cases} \Rightarrow a \rightarrow t$

$\begin{cases}
a \rightarrow t \\
a \rightarrow s \\
b = \text{lca}(s,t)
\end{cases} \Rightarrow a \rightarrow b$

Use a similar proof in Case 1, we would have an absurb situation.

As a result, the initial assumption that $\exists c \in V, c \rightarrow v, c \rightarrow t: c \nrightarrow b$ is false. In other words, $\forall c \in V, c\rightarrow v, c\rightarrow t: c \rightarrow b$

$\begin{cases}
b \rightarrow v \\
b \rightarrow t \\
\forall c \in V, c\rightarrow v, c\rightarrow t: c \rightarrow b
\end{cases} \Rightarrow b = \text{lca}(v,t)$

Because of the assumption, $f(u,v) > f(u,s)$

$\begin{cases}
u \rightarrow a \\
a \rightarrow v
\end{cases} \Rightarrow f(u,v) = f(u,a) + f(a,v)$

$\begin{cases}
u \rightarrow a \\
a \rightarrow s
\end{cases} \Rightarrow f(u,s) = f(u,a) + f(a,s)$

$\begin{cases}
f(u,v) > f(u,s) \\
f(u,v) = f(u,a) + f(a,v) \\
f(u,s) = f(u,a) + f(a,s)
\end{cases} \Rightarrow f(u,a) + f(a,v) > f(u,a) + f(a,s) \Rightarrow f(a,v) > f(a,s)$

$b = \text{lca}(v,t) \Rightarrow f(v,t) = f(v,b) + f(b,t)$

$\begin{cases}
b \rightarrow a \\
a \rightarrow v
\end{cases} \Rightarrow f(v,b) = f(v,a) + f(a,b)$

$\begin{cases}
b \rightarrow a \\
a \rightarrow s
\end{cases} \Rightarrow f(s,b) = f(s,a) + f(a,b)$

$f(a,v) > f(a,s) \Rightarrow f(v,a) + f(a,b) > f(s,a) + f(a,b) \Rightarrow f(v,b) > f(s,b)$

$\begin{cases}
f(v,b) > f(s,b) \\
f(v,t) = f(v,b) + f(b,t) \\
f(s,t) = f(s,b) + f(b,t)
\end{cases} \Rightarrow f(s,t) < f(v,t)$

$\Rightarrow$ the path from $s$ to $t$ is not a diameter of $T$, which is absurb

Because both cases lead to absurb situations, the initial assumption is false. In other words, a longest path from $u$ to $T$ must be a path from $u$ to $s$ or from $u$ to $t$.
</details>

## Part 2: Arrange the nodes based on $d_i$

We use the following theorem

Given an undirected tree $T = (V,E,w)$ with $w > 0$.

Let $L \subset V, L \ne \emptyset$ be a set of special nodes.

Let $d_i = \max_{j \in L} f(i,j)$

Choose $r \in V$ such that $d_r = \min_{i \in V} d_i$ and make it the root of the tree.

Prove that $\forall x,y \in V, x \rightarrow y: d_x \leq d_y$

<details>
<summary>Proof</summary>

First, we have the following lemma

$\forall x \in V, x \ne r: \forall y \in L, x \rightarrow y: d_x \ne f(x,y)$

Proof of the lemma:

Choose any $x, y$ such that $x \in V, x \ne r, y \in L, x \rightarrow y$

$\begin{cases}
\forall u \in V: d_r \leq d_u \\
x \in V
\end{cases} \Rightarrow d_x \leq d_u$

Assume, for a contradiction, that $d_x = f(x,y)$

Since $r$ is the root of $T$ and $x \in V$, $r \rightarrow x$

$\begin{cases}
r \rightarrow x \\
x \rightarrow y
\end{cases} \Rightarrow f(r,y) = f(r,x) + f(x,y)$

$\begin{cases}
y \in L \\
d_r = \max_{i \in V} f(r,i)
\end{cases} \Rightarrow d_r \geq f(r,y)$

Since $w > 0$ and $x \ne r$, $f(r,x) > 0 \Rightarrow f(r,x) + f(x,y) > f(x,y)$

$\begin{cases}
d_r \geq f(r,y) \\
f(r,y) = f(r,x) + f(x,y) \\
f(r,x) + f(x,y) > f(x,y) \\
f(x,y) = d_x
\end{cases} \Rightarrow d_r > d_x$, which contradicts with that $\forall u \in V: d_r \leq d_u$

As a result, the initial assumption is false. In other words, $d_x \ne f(x,y)$

Because of the way $x, y$ are chosen, I can conclude that

$\forall x \in V, x \ne r: \forall y \in L, x \rightarrow y: d_x \ne f(x,y)$

The lemma is true.

Let's go back to the original problem.

Choose any $x, y \in V$ such that $x \rightarrow y$

Suppose that $x = r$, then we have what we need to show.

Suppose that $x \ne r$, then by the above Lemma, we can show that $d_x = f(x,t)$ for some $t \in V, x \nrightarrow t$

Let $a = \text{lca}(x,t)$. Then, $\begin{cases} a \rightarrow x \\ a \rightarrow t \\ f(x,t) = f(x,a) + f(a,t) \end{cases}$

$\begin{cases}
a = \text{lca}(x,t) \\
x \nrightarrow t \\
x \rightarrow y
\end{cases} \Rightarrow a = \text{lca}(y,t) \Rightarrow f(y,t) = f(y,a) + f(a,t)$

$\begin{cases}
a \rightarrow x
x \rightarrow y
\end{cases} \Rightarrow f(y,a) = f(y,x) + f(x,a)$

We have: $f(y,t) = f(y,a) + f(a,t) = f(y,x) + f(x,a) + f(a,t) = f(y,x) + f(x,t) = f(y,x) + d_x$

Since $w > 0$, $f(y,x) \geq 0 \Rightarrow f(y,x) + d_x \geq d_x \Rightarrow f(y,t) \geq d_x$

Since $t \in L$ and $d_y = \max_{i \in L} f(y,i)$, $d_y \geq f(y,t)$

$\begin{cases}
d_y \geq f(y,t) \\
f(y,t) \geq d_x
\end{cases} \Rightarrow d_y \geq d_x$

By the way $x,y$ are chosen, I can conclude that

$\forall x,y \in V, x \rightarrow y: d_x \geq d_y$

which was what I need to show.
</details>

In this problem, $L = \{x| x \in V \text{ and } x \text{ is a leaf} \}$. We find the node $r$ with the lowest $d_r$, set it the root of the tree. The rooted tree would have the property $\forall x, y \in V, x \rightarrow y: d_x \leq d_y$

## Part 3: Get the set of nodes

Suppose that we have the largest set of nodes $S$ that we need to find. Then, one of the nodes with the lowest $d_i$ value would be the root of a subtree in our new rooted tree. Moreover, because $S$ is connected, other nodes in $S$ must belong to the subtree. We can also show that the set of nodes of $S$ is equal to the set of nodes $\{x | x \in \text{ (subtree with root } i) \text{ and } d_x \leq d_i + L\}$

So, for each query, we need to find the set $\{x | x \in \text{ (subtree with root } i) \text{ and } d_x \leq d_i + L\}$ for each node $i \in V$. We can do this in $O(n\log^2n)$ for each query using DFS:

To solve the problem for subtree with root $i$, solve the problem for each child of $i$. Each child of $i$ would gives $i$ a "list" of node. Combine these lists into the **largest** list, and remove the nodes that do not satisfy the condition $d_x \leq d_i + l$. Because the largest list is chosen, each node can only be moved to a new list $O(\log n)$ times (first, the node is in list with size 1, then size at least 2, then size at least 4, ...), and each node can only be removed once, so the total complexity without considering the time needed to add/remove a node from the list is $O(n\log n)$. To remove appropriate nodes fast, we need to use `pbds` to represent the lists, so the actual total complexity is $O(n \log^2 n)$

Since there are 50 queries, the final complexity is $O(qn\log^2 n)$. It takes ~0.5 seconds to run on Codeforces

My submission: https://codeforces.com/contest/516/submission/158614092