2
$\begingroup$

This is a question from a competitive programming competition. Given a tree with n nodes and a number k, find the total number of paths of length k in that tree.

I know for a fact that a solution can be found by using a centroid decomposition, however I am simply not sure on the exact relationship between the original tree and the tree we got from decomposition other than the fact that the length of path A -> B in the original tree can be decomposed into a sum A -> C + C -> B in the decomposed tree. Where C is the LCA(A, B). All distances being in the original tree.

$\endgroup$
7
$\begingroup$

This can be solved in $\mathcal{O}(n \log n)$ by using the smaller-to-larger merging technique. Root the tree at an arbitrary vertex. We will calculate for every subtree an array where the $d$th position indicates the number of nodes at depth $d$ in the subtree. Of course, the total size of these arrays could be $\mathcal{O}(n^{2})$, so we will not store them, instead building the array at every node by "scavenging" the arrays of its children.

To calculate the number of nodes at depth $d$ in the subtree of $i$, calculate the arrays for all children of $i$, then prepend zero to them, and sum them up, then finally add 1 to the number of nodes at depth $0$ (to count node $i$). For fast prepending, we store the arrays as dynamic arrays with inverted indexing. When we add two arrays together, we only care about the result, not about what happens to the orignal arrays, so we can just add the numbers in the smaller array to the numbers in the larger array and return the larger array.

Thus, if node $i$ has children $c_{1}, \dots, c_{m}$ and $T_{i}$ denotes the subtree of node $i$, we do an amount of work proportional to $w_{i} = |T_{i}| - \max_{j} |T_{c_{j}}|$. Therefore the total complexity is $\mathcal{O}(\sum_{i = 1}^{n} w_{i})$. Note that if $p$ is the parent of $i$ and $|T_{p}| < 2|T_{i}|$ then $|T_{i}|$ does not appear in the sum. Therefore, on any path from the root to some node $x$, at most $\log n$ sets we count appear, and those are exactly the sets $x$ appears in. Since every vertex appears in at most $\log n$ sets, their total size is $\mathcal{O}(n \log n)$. This is tight: consider the perfect binary tree of size $n = 2^{h} - 1$.

It remains to show how to calculate the number of paths of length $k$ from the depth information. To do this, when combining two arrays $A$ and $B$ at some node, just add $\sum_{i} A[i] B[k-i]$ to the answer. If $A$ is the sum of the arrays of previous children, and $B$ is the array of the current child we are handling, this corresponds to counting the paths with one endpoint in the current child's subtree, and the other in some earlier child's subtree. This can be done in $\mathcal{O}(\min(|A|, |B|))$ so it does not affect the runtime.

Here is a C++ implementation of the algorithm:

#include <iostream>
#include <vector>
#include <deque>
using namespace std;
using ll = long long;

ll combine(deque<int>& a, deque<int>& b, int k) {
    if (a.size() < b.size()) swap(a, b);
    int as = a.size();
    int bs = b.size();

    // Count paths of length k
    ll res = 0;
    for (int i = max(0, (k - bs) + 1); i < min(k + 1, as); ++i) {
        res += a[i] * b[k-i];
    }

    // Combine a and b
    for (int i = 0; i < bs; ++i) a[i] += b[i];
    return res;
}

pair<ll, deque<int>> solve(int i, int k, int p, const vector<vector<int>>& g) {
    pair<ll, deque<int>> res = {(k == 0), {1}};
    for (int t : g[i]) {
        if (t == p) continue;
        auto sub = solve(t, k, i, g);
        sub.second.push_front(0);
        res.first += sub.first + combine(res.second, sub.second, k);
    }
    return res;
}

int main() {
    int n, k;
    cin >> n >> k;
    vector<vector<int>> g(n);
    for (int i = 0; i < n-1; ++i) {
        int a, b;
        cin >> a >> b;
        --a; --b;
        g[a].push_back(b);
        g[b].push_back(a);
    }

    ll res = solve(0, k, -1, g).first;
    cout << res << '\n';
}
| cite | improve this answer | |
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.