CSES - Fixed-Length Paths I

Authors: Andi Qu, Benjamin Qi

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

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

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

Counting paths through uu

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

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

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

Applying centroid decomposition

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

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

Implementation

C++

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

Bonus Solution

It's also possible to solve this problem in linear time! See the solution for Fixed-Length Paths II for more details.

Join the USACO Forum!

Stuck on a problem, or don't understand a module? Join the USACO Forum and get help from other competitive programmers!