CSES - Fixed-Length Paths II

With Centroid Decomposition

Time Complexity: $\mathcal O(N \log^2 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.

This is a range sum, so we can use any range-sum data-structure (e.g. a BIT) to query and update $\texttt{cnt}$ efficiently. This adds an additional $\log N$ factor to the complexity.

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 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>;