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: 4Input: 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:
#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;
}
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);
}
}
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)
// 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)