CSES Solutions – Fixed-Length Paths II

Given a tree of n nodes, the task is to count the number of distinct paths that have at least k₁ and at most k₂ edges.

Example:

Input: n = 5, k1 = 2, k2 = 3, edges = {{1, 2}, {2 ,3}, {3, 4}, {3, 5}}
Output: 6

Input: n = 5, k1 = 1, k2 = 4, edges = { { 1, 2 }, { 2, 3 }, { 3, 4 }, { 3, 5 } }
Output: 10

Approach:

This solution is a simple extension of CSES Fixed-Length Paths I’s solution.

The main idea is to use centroid decomposition to divide the tree into smaller subtrees. Within each subtree, the get_cnt function counts paths based on the depth of nodes. The Fenwick tree is used to efficiently query and update the number of paths within specific depth ranges. The final answer is obtained by summing the counts from all subtrees.

Let’s break down the Ideas:

Centroid Decomposition:

• Find the centroid of the tree, which is the node that minimizes the maximum distance to any other node.
• Recursively decompose the tree into subtrees rooted at the centroids.

Fenwick Tree:

• Use a Fenwick tree (Binary Indexed Tree) to efficiently store and update counts of paths within the desired length range.

Counting Paths:

• For each subtree rooted at a centroid, count paths that pass through the centroid and satisfy the length range.
• Use the Fenwick tree to query the number of paths within the range for each depth level.

Steps-by-step approach:

• Perform centroid decomposition to break the tree into subtrees.
• Initialize the Fenwick tree with the root node.
• For each subtree, recursively count paths and update the Fenwick tree.
• Query the Fenwick tree to obtain the final count of paths within the specified length range.

Below is the implementation of the above approach:

C++

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

int n, a, b;

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

// Size of subtree rooted at each node
int subtree[200001];

// ans stores the final answer, bit is for Fenwick tree
// (Binary Indexed Tree)
ll ans = 0, bit[200001];

// Maximum depth encountered during centroid decomposition
int mx_depth;

// To mark nodes as processed during centroid decomposition
bool processed[200001];

// Function to calculate 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 a 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 Fenwick tree
void update(int pos, ll val)
{
    for (pos++; pos <= n; pos += pos & -pos)
        bit[pos] += val;
}

// Function to query Fenwick tree for sum in a range
ll query(int l, int r)
{
    ll ans = 0;
    for (r++; r; r -= r & -r)
        ans += bit[r];
    for (; l; l -= l & -l)
        ans -= bit[l];
    return ans;
}

// Function to count paths in a subtree within a certain
// depth range
void get_cnt(int node, int parent, bool filling,
             int depth = 1)
{

    // Depth exceeds limit, stop recursion
    if (depth > b)
        return;

    // Update maximum depth encountered
    mx_depth = max(mx_depth, depth);

    // Fill the Fenwick tree if filling is true
    if (filling)
        update(depth, 1);

    // Otherwise, query the Fenwick tree for counts
    else
        ans += query(max(0, a - depth), b - depth);
    for (int i : graph[node])
        if (!processed[i] && i != parent)
            get_cnt(i, node, filling, depth + 1);
}

// Centroid decomposition of the tree
void centroid_decomp(int node = 1)
{

    // Find centroid of the subtree
    int centroid
        = get_centroid(get_subtree_sizes(node) >> 1, node);

    // Mark centroid as processed
    processed[centroid] = true;

    // Initialize maximum depth encountered to 0
    mx_depth = 0;

    // Iterate through centroid's neighbors
    for (int i : graph[centroid])
        if (!processed[i]) {

            // Count paths passing through each neighbor
            get_cnt(i, centroid, false);

            // Update Fenwick tree for each neighbor
            get_cnt(i, centroid, true);
        }

    // Clear Fenwick tree values after processing
    for (int i = 1; i <= mx_depth; i++)
        update(i, -query(i, i));

    // Recursively decompose remaining subtrees
    for (int i : graph[centroid])
        if (!processed[i])
            centroid_decomp(i);
}

int main()
{
    // Static input
    n = 5; 
    a = 1; // or K1
    b = 4; // or K2
    vector<vector<int> > edges
        = { { 1, 2 }, { 2, 3 }, { 3, 4 }, { 3, 5 } };

    for (int i = 0; i < edges.size(); i++) {
        int u = edges[i][0];
        int v = edges[i][1];

        graph[u].push_back(v);
        graph[v].push_back(u);
    }

    // Initialize Fenwick tree with root node
    update(0, 1);

    // Perform centroid decomposition
    centroid_decomp();

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

Java

import java.util.*;

public class Main {
    static int n, a, b;
    static ArrayList<Integer>[] graph; // Adjacency list representation of the tree
    static int[] subtree; // Size of subtree rooted at each node
    static long ans = 0; // ans stores the final answer
    static long[] bit; // bit is for Fenwick tree (Binary Indexed Tree)
    static int mx_depth; // Maximum depth encountered during centroid decomposition
    static boolean[] processed; // To mark nodes as processed during centroid decomposition

    // Function to calculate 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 a 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 Fenwick tree
    static void update(int pos, long val) {
        for (pos++; pos <= n; pos += pos & -pos) {
            bit[pos] += val;
        }
    }

    // Function to query Fenwick tree for sum in a range
    static long query(int l, int r) {
        long ans = 0;
        for (r++; r > 0; r -= r & -r) {
            ans += bit[r];
        }
        for (; l > 0; l -= l & -l) {
            ans -= bit[l];
        }
        return ans;
    }

    // Function to count paths in a subtree within a certain depth range
    static void getCnt(int node, int parent, boolean filling, int depth) {
        if (depth > b) return;
        mx_depth = Math.max(mx_depth, depth);
        if (filling) update(depth, 1);
        else ans += query(Math.max(0, a - depth), b - depth);
        for (int i : graph[node]) {
            if (!processed[i] && i != parent) {
                getCnt(i, node, filling, depth + 1);
            }
        }
    }

    // Centroid decomposition of the tree
    static void centroidDecomp(int node) {
        int centroid = getCentroid(getSubtreeSizes(node, 0) >> 1, node, 0);
        processed[centroid] = true;
        mx_depth = 0;
        for (int i : graph[centroid]) {
            if (!processed[i]) {
                getCnt(i, centroid, false, 1);
                getCnt(i, centroid, true, 1);
            }
        }
        for (int i = 1; i <= mx_depth; i++) {
            update(i, -query(i, i));
        }
        for (int i : graph[centroid]) {
            if (!processed[i]) {
                centroidDecomp(i);
            }
        }
    }

    public static void main(String[] args) {
        n = 5;
        a = 1;
        b = 4;
        ArrayList<ArrayList<Integer>> edges = new ArrayList<>();
        edges.add(new ArrayList<>(Arrays.asList(1, 2)));
        edges.add(new ArrayList<>(Arrays.asList(2, 3)));
        edges.add(new ArrayList<>(Arrays.asList(3, 4)));
        edges.add(new ArrayList<>(Arrays.asList(3, 5)));

        graph = new ArrayList[n + 1];
        for (int i = 1; i <= n; i++) {
            graph[i] = new ArrayList<>();
        }

        for (ArrayList<Integer> edge : edges) {
            int u = edge.get(0);
            int v = edge.get(1);
            graph[u].add(v);
            graph[v].add(u);
        }

        subtree = new int[n + 1];
        processed = new boolean[n + 1];
        bit = new long[n + 1];

        update(0, 1);
        centroidDecomp(1);

        System.out.println(ans); // Output the final answer
    }
}
//This code is contributed by Utkarsh.

Python3

import sys

# Adjacency list representation of the tree
graph = [[] for _ in range(200001)]

# Size of subtree rooted at each node
subtree = [0] * 200001

# ans stores the final answer, bit is for Fenwick tree
bit = [0] * 200001

# Maximum depth encountered during centroid decomposition
mx_depth = 0

# To mark nodes as processed during centroid decomposition
processed = [False] * 200001

# Function to calculate subtree sizes rooted at each node
def get_subtree_sizes(node, parent=0):
    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 a subtree
def get_centroid(desired, node, parent=0):
    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 Fenwick tree
def update(pos, val):
    while pos <= n:
        bit[pos] += val
        pos += pos & -pos

# Function to query Fenwick tree for sum in a range
def query(l, r):
    ans = 0
    while r:
        ans += bit[r]
        r -= r & -r
    while l:
        ans -= bit[l]
        l -= l & -l
    return ans

# Function to count paths in a subtree within a certain depth range
def get_cnt(node, parent, filling, depth=1):
    global mx_depth, ans
    if depth > b:
        return
    mx_depth = max(mx_depth, depth)
    if filling:
        update(depth, 1)
    else:
        ans += query(max(0, a - depth), b - depth)
    for i in graph[node]:
        if not processed[i] and i != parent:
            get_cnt(i, node, filling, depth + 1)

# Centroid decomposition of the tree
def centroid_decomp(node=1):
    global mx_depth
    centroid = get_centroid(get_subtree_sizes(node) >> 1, node)
    processed[centroid] = True
    mx_depth = 0
    for i in graph[centroid]:
        if not processed[i]:
            get_cnt(i, centroid, False)
            get_cnt(i, centroid, True)
    for i in range(1, mx_depth + 1):
        update(i, -query(i, i))
    for i in graph[centroid]:
        if not processed[i]:
            centroid_decomp(i)

# Static input
n = 5
a = 1
b = 4
edges = [[1, 2], [2, 3], [3, 4], [3, 5]]

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

# Final answer
ans = 0
centroid_decomp()
print(ans)

# This code is contributed by Ayush Mishra

JavaScript

// Initialize global variables
let n, a, b;
let graph = new Array(200001);
let subtree = new Array(200001).fill(0);
let ans = 0;
let bit = new Array(200001).fill(0);
let mx_depth;
let processed = new Array(200001).fill(false);

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

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

// Function to update Fenwick tree
function update(pos, val) {
    for (pos++; pos <= n; pos += pos & -pos) {
        bit[pos] += val;
    }
}

// Function to query Fenwick tree for sum in a range
function query(l, r) {
    let ans = 0;
    for (r++; r; r -= r & -r) {
        ans += bit[r];
    }
    for (; l; l -= l & -l) {
        ans -= bit[l];
    }
    return ans;
}

// Function to count paths in a subtree within a certain depth range
function get_cnt(node, parent, filling, depth = 1) {
    if (depth > b) {
        return;
    }
    mx_depth = Math.max(mx_depth, depth);
    if (filling) {
        update(depth, 1);
    } else {
        ans += query(Math.max(0, a - depth), b - depth);
    }
    for (let i of graph[node]) {
        if (!processed[i] && i !== parent) {
            get_cnt(i, node, filling, depth + 1);
        }
    }
}

// Centroid decomposition of the tree
function centroid_decomp(node = 1) {
    let centroid = get_centroid(get_subtree_sizes(node) >> 1, node);
    processed[centroid] = true;
    mx_depth = 0;
    for (let i of graph[centroid]) {
        if (!processed[i]) {
            get_cnt(i, centroid, false);
            get_cnt(i, centroid, true);
        }
    }
    for (let i = 1; i <= mx_depth; i++) {
        update(i, -query(i, i));
    }
    for (let i of graph[centroid]) {
        if (!processed[i]) {
            centroid_decomp(i);
        }
    }
}

// Main function
function main() {
    // Static input
    n = 5;
    a = 1; // or K1
    b = 4; // or K2
    let edges = [[1, 2], [2, 3], [3, 4], [3, 5]];

    for (let i = 1; i <= n; i++) {
        graph[i] = [];
    }

    for (let i = 0; i < edges.length; i++) {
        let u = edges[i][0];
        let v = edges[i][1];

        graph[u].push(v);
        graph[v].push(u);
    }

    // Initialize Fenwick tree with root node
    update(0, 1);

    // Perform centroid decomposition
    centroid_decomp();

    // Output the final answer
    console.log(ans);
}

// Invoke the main function
main();

Output

10

Time complexity: O(N log N)
Auxiliary space: O(N) + O(log N)