Описание задачи 📋
Дано неориентированное дерево с nn узлами, пронумерованными от 0 до n−1. Также даны:
- массив рёбер edges, где каждое ребро связывает два узла,
- массив значений values, где значение values[i] связано с узлом i,
- число k.
- гарантируется, что сумма всех значений values кратна k
Необходимо найти максимальное количество компонентов, на которые можно разделить дерево, удаляя рёбра, чтобы сумма значений в каждом компоненте делилась на k.
Идея 💡
Будем разбивать дерево на компоненты рекурсивно (DFS), начиная с произвольного узла.
Несложно показать, что родительская вершина должна быть в одной компененте со всеми дочерними, для которых сумма значений в соответствующем поддереве не делится на k.
Детали подхода 🚀
-
Построение графа:
- Представляем дерево в виде списка смежности для удобного обхода.
-
Обход дерева:
- Запускаем DFS от корня.
- Для каждого узла вычисляем остаток от суммы поддерева (
sum%k) и проверяем, можно ли выделить текущую часть как отдельный компонент.
- Рекурсивно собираем остатки и компоненты из всех дочерних узлов.
-
Подсчёт компонентов:
- Если сумма узлов поддерева делится на
k, увеличиваем счётчик компонентов.
- Возвращаем количество валидных компонентов в дереве.
Асимптотика 📊
- Временная сложность:
O(n)
- Каждый узел и ребро обрабатывается один раз.
- Пространственная сложность:
O(n)
- Используется память для хранения графа и стека вызовов DFS.
Код 💻
impl Solution {
pub fn max_k_divisible_components(n: i32, edges: Vec<Vec<i32>>, values: Vec<i32>, k: i32) -> i32 {
// Build adjacency list
let mut graph = vec![vec![]; n as usize];
for edge in edges {
let (u, v) = (edge[0] as usize, edge[1] as usize);
graph[u].push(v);
graph[v].push(u);
}
// Perform DFS starting from node 0
let mut visited = vec![false; n as usize];
let (total_sum, total_components) = Self::split_components(0, &graph, &values, k, &mut visited);
// Validate the total sum
assert!(total_sum == 0, "Invalid tree: total values sum is not divisible by k");
total_components
}
fn split_components(
node: usize,
graph: &Vec<Vec<usize>>,
values: &Vec<i32>,
k: i32,
visited: &mut Vec<bool>,
) -> (i32, i32) {
visited[node] = true;
let mut current_sum = values[node] % k;
let mut current_components = 0;
for &neighbor in &graph[node] {
if !visited[neighbor] {
let (child_sum, child_components) = Self::split_components(neighbor, graph, values, k, visited);
current_sum = (current_sum + child_sum) % k;
current_components += child_components;
}
}
if current_sum == 0 {
current_components += 1;
}
(current_sum, current_components)
}
}