Solution to Codeforces 2040E | Yuhang's Blog

Link to problem.

We build the tree so that vertex $1$ is the root and each other vertex $v$ has parent $\text{par}(v)$. Let $f(v, p, i)$ be the expected number of steps to reach the root if we start from vertex $v$ with $p$ coins, and the index of the next step has parity $i$, represented by $0$ or $1$.

We now try to find the recurrence relation for $f$. The base case is trivial: for $v = 1$, we have $f(1, p, i) =0$. Now suppose $v \neq 1$.

When $i$ is $1$, this case is fairly easy. We have
$$
f(v, p, 1) = f(\text{par}(v), p, 0) + 1,
$$
as we can only move one step to $v$’s parent.

Now suppose $i$ is $0$. Suppose $p = 0$, so that we have no coins to spend. In this case, we will be forced to move uniformly at random to an adjacent vertex. An adjacent vertex is either the parent of $v$ or a child of $v$. Suppose $v$ has in total $s$ adjacent vertices (and thus $s-1$ children). Then the probability is $\frac 1s$ to reach $v$’s parent in one step. If, however, we choose any child of $v$, then the next step is an odd step and we are forced to go back to $v$, so we are back to the initial state with two more steps spent.

This fits into something like a negative binomial distribution: basically, we are just repeatedly trying until $v$’s parent is chosen in this random process. In other words, we have probability $\frac 1s$ to reach $v$’s parent in one step; probability $\frac{s-1}s \frac{1}{s}$ to reach $v$’s parent in three steps; probability $(\frac{s-1}s)^2 \frac{1}{s}$ in five steps; and so on. The expected number of steps it takes to move to $v$’s parent is thus the following series:

$$
\displaylines{
\frac1s + 3\left(\frac{s-1}s\right) \frac{1}{s} + 5\left(\frac{s-1}s\right)^2 \frac{1} {s} + … \\ = \frac{1}{s} \sum_{k \ge 0} (2 k + 1) \left(\frac{s-1}s\right)^k = \frac{1}{s} \sum_{k \ge 0} (2 k + 1) x^k
}
$$
where we denote $x = \frac{s-1}{s}$. We know this following series for $-1<x<1$:
$$
\sum_{k\ge 0} x^k = \frac{1}{1-x}.
$$
Taking the derivative w.r.t $x$, we have
$$
\sum_{k\ge 0} k x^{k-1} = \frac{1}{(1-x)^2}.
$$
Hence
$$
\sum_{k\ge 0} 2k x^{k} = \frac{2x}{(1-x)^2}.
$$
Therefore
$$
\frac{1}{s} \sum_{k \ge 0} (2 k + 1) x^k = \frac 1s \left( \frac{2x}{(1-x)^2} + \frac{1}{1-x}\right) = 2s-1.
$$
So we have $f(v, 0, 0) = f(\text{par}(v), 0, 1) + 2s - 1$.

When $p \neq 0$, we have the chance to spend a coin and move to $v$’s parent directly, so the result is $f(v, p, 0) = \min(f(\text{par}(v), p - 1, 1) + 1, f(\text{par}(v), p, 1) + 2s - 1)$.

Further remark. Problems like this with a random process on a graph (or a tree) usually creates circular dependencies in the recurrence relations and one would expect to use Gaussian Elimination or otherwise to solve a system of linear equations stipulated by the recurrence relations. This problem is however a bit different; in fact, Gaussian Elimination which has time complexity $O(N^3)$ cannot work on the scale of the input $2\times 10^3$ of this problem, and we have to carefully exploit the properties specific to this problem.

The C++ code follows.

#include<bits/stdc++.h>
using namespace std;
#define rep(i,from,to) for(int i=(int)(from);i<=(int)(to);++i)
#define rev(i,from,to) for(int i=(int)(from);i>=(int)(to);--i)
#define endl '\n'
#define mmst(a,x) memset(a, x, sizeof(a))
typedef long long ll; 

const ll MOD = 998244353;
const int N = 2123;
vector<int> g[N];
int fa[N];
int n, q;


void build(int u) {
    for (auto v: g[u]) if (v != fa[u]) {
        fa[v] = u;
        build(v);
    }
}

ll mem[N][N][2];
ll dp(ll v, ll p, ll i) {
    if (v == 1) return 0; 
    if (mem[v][p][i] != -1) return mem[v][p][i];
    if (i == 1) {
        return mem[v][p][i] = 1ll + dp(fa[v], p, 0);
    } else {
        ll sz = g[v].size();
        if (p == 0) {
            return mem[v][p][i] = ((2 * sz - 1) + dp(fa[v], p, 1));
        }
        return mem[v][p][i] = min((2 * sz - 1) + dp(fa[v], p, 1), dp(fa[v], p - 1, 1) + 1);
    }
}


void solve() {
    cin >> n >> q;
    rep(i,1,n) g[i].clear(), fa[i] = 0;
    rep(i,1,n-1) {
        int a, b;
        cin >> a >> b; 
        g[a].push_back(b);
        g[b].push_back(a);
    }

    build(1);
    rep(v,0,n) rep(p,0,n) rep(i,0,1) mem[v][p][i] = -1;
    rep(_, 1, q) {
        int v, p;
        rin >> v >> p;
        cout << (dp(v, p, 1) % MOD) << endl;
    }
}

int main() {
    int T; 
    cin >> T; 
    while (T--) {
        solve();
    }
    return 0;
}