CSES - Fixed-Length Paths I

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

Given a node $u$, we know that any path in the tree either passes $u$ or is completely contained in one of $u$'s children's subtrees. This means that to solve the problem on a tree containing $u$, we can:

1. Count the number of paths of length $k$ that pass through $u$.
2. Remove $u$ and its incident edges from the tree.
3. Recursively solve the problem on the new trees that are formed.

Counting paths through $u$

We can count the number of paths passing through $u$ in a tree of size $M$ in $\mathcal O(M)$ time.

We process each of $u$'s children's subtrees in order. Let $\texttt{cnt}_i[d]$ be the number of nodes we've visited with depth $d$ after processing the first $i$ of $u$'s children's subtrees. A node with depth $d$ in $u$'s $i$-th child's subtree will contribute $\texttt{cnt}_{i - 1}[k - d]$ paths to the answer.

Using two DFSes for each child, we can calculate its contribution and then update $\texttt{cnt}$ in $\mathcal O(M)$ time total.

Applying centroid decomposition

We can combine the previous idea with centroid decomposition to solve the whole problem in $\mathcal O(N \log N)$ time.

When processing some tree, simply let the $u$ we use be that tree's centroid. Since the centroid tree has a depth of $\mathcal O(\log N)$ and each layer in the centroid tree takes amortized $\mathcal O(N)$ time to process, the total complexity is $\mathcal O(N log N)$.

Implementation

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

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

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

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