CSES Solutions – Fixed-Length Paths I

Given a tree of n nodes, your task is to count the number of distinct paths that consist of exactly k edges.

Example:

Input: n = 5 , k = 2, edges = {{1, 2}, {2, 3}, {3, 4}, {3, 5}
Output: 4

Input: n = 5 , k = 3, edges = {{1, 2}, {2, 3}, {3, 4}, {3, 5}
Output: 2

Approach:

The main idea is that we can use centroid decomposition to efficiently count the number of paths of length k in a tree. Centroid decomposition allows us to decompose a tree into a set of smaller subtrees, each of which is rooted at a centroid. The centroid of a tree is the node that minimizes the maximum distance to any other node in the tree.

Once we have decomposed the tree into subtrees, we can count the number of paths of length k in each subtree. We can do this by processing each subtree in a bottom-up manner, starting from the leaves. For each node in the subtree, we count the number of paths of length k that pass through that node. We can do this by considering all of the paths of length k-1 that end at the node’s children.

Steps-by-step approach:

• Perform centroid decomposition on the tree.
• For each subtree, process the nodes in a bottom-up manner, starting from the leaves.
• For each node in the subtree, count the number of paths of length k that pass through that node.
• Add the number of paths of length k that pass through each node in the subtree to the total count.

Below is the implementation of the above approach:

C++

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

int n, k;

// Adjacency list representation of the graph
vector<int> graph[200001];

// Stores subtree sizes for each node
int subtree[200001];

// Final answer to be computed
ll ans = 0;

// Counter array to count the number of nodes at each depth
// in the subtree, and mx_depth to store the maximum depth
// encountered
int cnt[200001]{ 1 }, mx_depth;

// Keeps track of processed nodes during centroid
// decomposition
bool processed[200001];

// Function to compute subtree sizes rooted at each node
int get_subtree_sizes(int node, int parent = 0)
{
    subtree[node] = 1;
    for (int i : graph[node])
        if (!processed[i] && i != parent)
            subtree[node] += get_subtree_sizes(i, node);
    return subtree[node];
}

// Function to find the centroid of the subtree
int get_centroid(int desired, int node, int parent = 0)
{
    for (int i : graph[node])
        if (!processed[i] && i != parent
            && subtree[i] >= desired)
            return get_centroid(desired, i, node);
    return node;
}

// Function to update the counter array with node counts at
// each depth
void get_cnt(int node, int parent, bool filling,
             int depth = 1)
{
    if (depth > k)
        return;
    mx_depth = max(mx_depth, depth);
    if (filling)
        cnt[depth]++;
    else
        ans += cnt[k - depth];
    for (int i : graph[node])
        if (!processed[i] && i != parent)
            get_cnt(i, node, filling, depth + 1);
}

// Centroid decomposition function
void centroid_decomp(int node = 1)
{
    // Finding the centroid of the subtree rooted at 'node'
    int centroid
        = get_centroid(get_subtree_sizes(node) >> 1, node);

    // Marking the centroid as processed
    processed[centroid] = true;

    // Resetting the maximum depth encountered
    mx_depth = 0;

    // Updating counters for each child of the centroid
    for (int i : graph[centroid])
        if (!processed[i]) {
            get_cnt(i, centroid, false);
            get_cnt(i, centroid, true);
        }
    // Resetting counter array for further iterations
    fill(cnt + 1, cnt + mx_depth + 1, 0);
    // Recursively decomposing the subtrees rooted at the
    // children of the centroid
    for (int i : graph[centroid])
        if (!processed[i])
            centroid_decomp(i);
}

// Driver code
int main()
{
    cin.tie(0)->sync_with_stdio(0);
    // Static input
    n = 5; // Number of nodes
    k = 2; // Maximum depth
    // Static graph edges
    graph[1].push_back(2);
    graph[2].push_back(1);
    graph[2].push_back(3);
    graph[3].push_back(2);
    graph[3].push_back(4);
    graph[3].push_back(5);
    graph[4].push_back(3);
    graph[5].push_back(3);

    // Performing centroid decomposition and computing the
    // final answer
    centroid_decomp();

    // Outputting the final answer
    cout << ans;
    return 0;
}

Java

import java.util.ArrayList;

public class Main {

    static int n, k;

    // Adjacency list representation of the graph
    static ArrayList<Integer>[] graph
        = new ArrayList[200001];

    // Stores subtree sizes for each node
    static int[] subtree = new int[200001];

    // Final answer to be computed
    static long ans = 0;

    // Counter array to count the number of nodes at each
    // depth in the subtree, and mx_depth to store the
    // maximum depth encountered
    static int[] cnt = new int[200001];
    static int mx_depth;

    // Keeps track of processed nodes during centroid
    // decomposition
    static boolean[] processed = new boolean[200001];

    // Function to compute subtree sizes rooted at each node
    static int getSubtreeSizes(int node, int parent)
    {
        subtree[node] = 1;
        for (int i : graph[node])
            if (!processed[i] && i != parent)
                subtree[node] += getSubtreeSizes(i, node);
        return subtree[node];
    }

    // Function to find the centroid of the subtree
    static int getCentroid(int desired, int node,
                           int parent)
    {
        for (int i : graph[node])
            if (!processed[i] && i != parent
                && subtree[i] >= desired)
                return getCentroid(desired, i, node);
        return node;
    }

    // Function to update the counter array with node counts
    // at each depth
    static void getCnt(int node, int parent,
                       boolean filling, int depth)
    {
        if (depth > k)
            return;
        mx_depth = Math.max(mx_depth, depth);
        if (filling)
            cnt[depth]++;
        else
            ans += cnt[k - depth];
        for (int i : graph[node])
            if (!processed[i] && i != parent)
                getCnt(i, node, filling, depth + 1);
    }

    // Centroid decomposition function
    static void centroidDecomp(int node)
    {
        // Finding the centroid of the subtree rooted at
        // 'node'
        int centroid = getCentroid(
            getSubtreeSizes(node, 0) >> 1, node, 0);

        // Marking the centroid as processed
        processed[centroid] = true;

        // Resetting the maximum depth encountered
        mx_depth = 0;

        // Updating counters for each child of the centroid
        for (int i : graph[centroid])
            if (!processed[i]) {
                getCnt(i, centroid, false, 1);
                getCnt(i, centroid, true, 1);
            }

        // Resetting counter array for further iterations
        for (int i = 1; i <= mx_depth; i++)
            cnt[i] = 0;

        // Recursively decomposing the subtrees rooted at
        // the children of the centroid
        for (int i : graph[centroid])
            if (!processed[i])
                centroidDecomp(i);
    }

    // Driver code
    public static void main(String[] args)
    {
        // Static input
        n = 5; // Number of nodes
        k = 2; // Maximum depth

        // Initializing adjacency list
        for (int i = 0; i <= n; i++)
            graph[i] = new ArrayList<>();

        // Static graph edges
        graph[1].add(2);
        graph[2].add(1);
        graph[2].add(3);
        graph[3].add(2);
        graph[3].add(4);
        graph[3].add(5);
        graph[4].add(3);
        graph[5].add(3);

        // Performing centroid decomposition and computing
        // the final answer
        centroidDecomp(1);

        // Outputting the final answer
        System.out.println(ans);
    }
}

Python

from collections import defaultdict

# Adjacency list representation of the graph
graph = defaultdict(list)

# Stores subtree sizes for each node
subtree = [0] * 200001

# Final answer to be computed
ans = 0

# Counter array to count the number of nodes at each depth
# in the subtree, and mx_depth to store the maximum depth
# encountered
cnt = [1] + [0] * 200000
mx_depth = 0

# Keeps track of processed nodes during centroid
# decomposition
processed = [False] * 200001

# Function to compute subtree sizes rooted at each node


def get_subtree_sizes(node, parent):
    subtree[node] = 1
    for i in graph[node]:
        if not processed[i] and i != parent:
            subtree[node] += get_subtree_sizes(i, node)
    return subtree[node]

# Function to find the centroid of the subtree


def get_centroid(desired, node, parent):
    for i in graph[node]:
        if not processed[i] and i != parent and subtree[i] >= desired:
            return get_centroid(desired, i, node)
    return node

# Function to update the counter array with node counts at
# each depth


def get_cnt(node, parent, filling, depth):
    global mx_depth, ans
    if depth > k:
        return
    mx_depth = max(mx_depth, depth)
    if filling:
        cnt[depth] += 1
    else:
        ans += cnt[k - depth]
    for i in graph[node]:
        if not processed[i] and i != parent:
            get_cnt(i, node, filling, depth + 1)

# Centroid decomposition function


def centroid_decomp(node):
    global ans, mx_depth
    # Finding the centroid of the subtree rooted at 'node'
    centroid = get_centroid(get_subtree_sizes(node, 0) >> 1, node, 0)

    # Marking the centroid as processed
    processed[centroid] = True

    # Resetting the maximum depth encountered
    mx_depth = 0

    # Updating counters for each child of the centroid
    for i in graph[centroid]:
        if not processed[i]:
            get_cnt(i, centroid, False, 1)
            get_cnt(i, centroid, True, 1)

    # Resetting counter array for further iterations
    cnt[1:mx_depth+1] = [0] * (mx_depth)

    # Recursively decomposing the subtrees rooted at the
    # children of the centroid
    for i in graph[centroid]:
        if not processed[i]:
            centroid_decomp(i)


# Static input
n = 5  # Number of nodes
k = 2  # Maximum depth

# Static graph edges
edges = [
    (1, 2),
    (2, 3),
    (3, 4),
    (3, 5)
]

for u, v in edges:
    graph[u].append(v)
    graph[v].append(u)

# Performing centroid decomposition and computing the
# final answer
centroid_decomp(1)

# Outputting the final answer
print(ans)

JavaScript

// Function to compute subtree sizes rooted at each node
function getSubtreeSizes(node, parent, graph, processed, subtree) {
    subtree[node] = 1;
    for (let i of graph[node]) {
        if (!processed[i] && i !== parent) {
            subtree[node] += getSubtreeSizes(i, node, graph, processed, subtree);
        }
    }
    return subtree[node];
}

// Function to find the centroid of the subtree
function getCentroid(desired, node, parent, graph, processed, subtree) {
    for (let i of graph[node]) {
        if (!processed[i] && i !== parent && subtree[i] >= desired) {
            return getCentroid(desired, i, node, graph, processed, subtree);
        }
    }
    return node;
}

// Function to update the counter array with node counts at each depth
function getCount(node, parent, filling, depth, graph, processed, cnt, ans, k) {
    if (depth > k) return;
    mxDepth = Math.max(mxDepth, depth);
    if (filling) cnt[depth]++;
    else ans[0] += cnt[k - depth];
    for (let i of graph[node]) {
        if (!processed[i] && i !== parent) {
            getCount(i, node, filling, depth + 1, graph, processed, cnt, ans, k);
        }
    }
}

// Centroid decomposition function
function centroidDecomp(node, graph, processed, subtree, cnt, ans, k) {
    let centroid = getCentroid(getSubtreeSizes(node, 0, graph, processed, subtree) >> 1, node, 0, graph, processed, subtree);
    processed[centroid] = true;
    mxDepth = 0;
    for (let i of graph[centroid]) {
        if (!processed[i]) {
            getCount(i, centroid, false, 1, graph, processed, cnt, ans, k);
            getCount(i, centroid, true, 1, graph, processed, cnt, ans, k);
        }
    }
    cnt.fill(0, 1, mxDepth + 1);
    for (let i of graph[centroid]) {
        if (!processed[i]) {
            centroidDecomp(i, graph, processed, subtree, cnt, ans, k);
        }
    }
}

// Driver code
function main() {
    // Static input
    let n = 5; // Number of nodes
    let k = 2; // Maximum depth
    let ans = [0]; // Final answer
    let graph = new Array(n + 1).fill(null).map(() => []);
    let processed = new Array(n + 1).fill(false);
    let subtree = new Array(n + 1).fill(0);
    let cnt = new Array(n + 1).fill(1);
    
    // Static graph edges
    graph[1].push(2);
    graph[2].push(1);
    graph[2].push(3);
    graph[3].push(2);
    graph[3].push(4);
    graph[3].push(5);
    graph[4].push(3);
    graph[5].push(3);

    // Performing centroid decomposition and computing the final answer
    centroidDecomp(1, graph, processed, subtree, cnt, ans, k);

    // Outputting the final answer
    console.log(ans[0]);
}

main();

// This code is contributed by Ayush Mishra

Output

4

Time Complexity: O(n log n)
Auxiliary Space: O(n + m)