CSES - Fixed-Length Paths II

With Centroid Decomposition

Time Complexity: $\mathcal O(N \log N)$

This solution is a simple extension of CSES Fixed-Length Paths I's solution.

Since we want to count the number of paths of length between $a$ and $b$ instead of a fixed length $k$, a node with depth $d$ in $u$'s $i$-th child's subtree will contribute $\sum_{x = a - d}^{b - d}\texttt{cnt}_{i - 1}[x]$ paths instead of $\texttt{cnt}_{i - 1}[k - d]$ paths to the answer.

We could calculate this is using a range-sum data-structure (e.g. a BIT). However this would add an additional $\log N$ factor to the complexity which might be too slow.

Another way is to calculate and maintain the partial sums of $\texttt{cnt}$. We can calculate the partial sum in time proportional to the size of the subtree, and each time $d$ is incremented we can update the partial sum in $\mathcal O(1)$ time.

This allows us to keep the total work done at each level of the centroid decomposition at $\mathcal O(N)$ and keep the overall time complexity at $\mathcal O(N \log N)$.

Copy
#include <bits/stdc++.h>
typedef long long ll;
using namespace std;

int n, a, b;
vector<int> graph[200001];
int subtree[200001];

ll ans = 0, bit[200001];
int total_cnt[200001]{1}, mx_depth;

Faster Solution

Time Complexity: $\mathcal O(N)$

We use the same idea as the above solution but with prefix sums instead of a BIT and with small-to-large merging in a single DFS instead of centroid decomposition. This is because we can merge two prefix sum arrays $A$ and $B$ in $\mathcal O(\min(|A|, |B|))$ time.

If we let $d_i$ denote the maximum depth in node $i$'s subtree, then the number of operations $S$ that our algorithm performs can be calculated by

Now,

and

Summing these two expressions and adding $N$ gives

which is $\mathcal O(N)$. Below is Ben's code implementing this idea:

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

using ll = long long;
using db = long double;  // or double, if TL is tight
using str = string;      // yay python!

using pi = pair<int, int>;
using pl = pair<ll, ll>;
using pd = pair<db, db>;