Skip to content

2872. Maximum Number of K-Divisible Components

SourceBiweekly Contest 114 Q4DifficultyHardRating1967

Description

There is an undirected tree with n nodes labeled from 0 to n - 1. You are given the integer n and a 2D integer array edges of length n - 1, where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

You are also given a 0-indexed integer array values of length n, where values[i] is the value associated with the ith node, and an integer k.

A valid split of the tree is obtained by removing any set of edges, possibly empty, from the tree such that the resulting components all have values that are divisible by k, where the value of a connected component is the sum of the values of its nodes.

Return the maximum number of components in any valid split.

 

Example 1:

Input: n = 5, edges = [[0,2],[1,2],[1,3],[2,4]], values = [1,8,1,4,4], k = 6
Output: 2
Explanation: We remove the edge connecting node 1 with 2. The resulting split is valid because:
- The value of the component containing nodes 1 and 3 is values[1] + values[3] = 12.
- The value of the component containing nodes 0, 2, and 4 is values[0] + values[2] + values[4] = 6.
It can be shown that no other valid split has more than 2 connected components.

Example 2:

Input: n = 7, edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]], values = [3,0,6,1,5,2,1], k = 3
Output: 3
Explanation: We remove the edge connecting node 0 with 2, and the edge connecting node 0 with 1. The resulting split is valid because:
- The value of the component containing node 0 is values[0] = 3.
- The value of the component containing nodes 2, 5, and 6 is values[2] + values[5] + values[6] = 9.
- The value of the component containing nodes 1, 3, and 4 is values[1] + values[3] + values[4] = 6.
It can be shown that no other valid split has more than 3 connected components.

 

Constraints:

  • 1 <= n <= 3 * 104
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • values.length == n
  • 0 <= values[i] <= 109
  • 1 <= k <= 109
  • Sum of values is divisible by k.
  • The input is generated such that edges represents a valid tree.

Solutions

Solution 1: DFS

Thinking

The whole tree sum is divisible by \(k\), so deleting a subtree whose sum is also divisible by \(k\) leaves valid components. A bottom-up DFS accumulates subtree sums and counts every subtree whose sum is \(0\) modulo \(k\).

We note that the problem guarantees the sum of all node values in the entire tree is divisible by \(k\). Therefore, if we remove a subtree whose sum of elements is divisible by \(k\), the sum of node values in each of the remaining connected components must also be divisible by \(k\).

Thus, we can use a depth-first search approach, starting from the root node to traverse the entire tree. For each node, we calculate the sum of all node values in its subtree. If this sum is divisible by \(k\), we increment the answer by one.

The time complexity is \(O(n)\) and the space complexity is \(O(n)\), where \(n\) is the number of nodes in the tree.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution:
    def maxKDivisibleComponents(
        self, n: int, edges: List[List[int]], values: List[int], k: int
    ) -> int:
        def dfs(i: int, fa: int) -> int:
            s = values[i]
            for j in g[i]:
                if j != fa:
                    s += dfs(j, i)
            nonlocal ans
            ans += s % k == 0
            return s

        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        ans = 0
        dfs(0, -1)
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Solution {
    private int ans;
    private List<Integer>[] g;
    private int[] values;
    private int k;

    public int maxKDivisibleComponents(int n, int[][] edges, int[] values, int k) {
        g = new List[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int[] e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        this.values = values;
        this.k = k;
        dfs(0, -1);
        return ans;
    }

    private long dfs(int i, int fa) {
        long s = values[i];
        for (int j : g[i]) {
            if (j != fa) {
                s += dfs(j, i);
            }
        }
        ans += s % k == 0 ? 1 : 0;
        return s;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
public:
    int maxKDivisibleComponents(int n, vector<vector<int>>& edges, vector<int>& values, int k) {
        int ans = 0;
        vector<int> g[n];
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        auto dfs = [&](this auto&& dfs, int i, int fa) -> long long {
            long long s = values[i];
            for (int j : g[i]) {
                if (j != fa) {
                    s += dfs(j, i);
                }
            }
            ans += s % k == 0;
            return s;
        };
        dfs(0, -1);
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
func maxKDivisibleComponents(n int, edges [][]int, values []int, k int) (ans int) {
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    var dfs func(int, int) int
    dfs = func(i, fa int) int {
        s := values[i]
        for _, j := range g[i] {
            if j != fa {
                s += dfs(j, i)
            }
        }
        if s%k == 0 {
            ans++
        }
        return s
    }
    dfs(0, -1)
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
function maxKDivisibleComponents(
    n: number,
    edges: number[][],
    values: number[],
    k: number,
): number {
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    let ans = 0;
    const dfs = (i: number, fa: number): number => {
        let s = values[i];
        for (const j of g[i]) {
            if (j !== fa) {
                s += dfs(j, i);
            }
        }
        if (s % k === 0) {
            ++ans;
        }
        return s;
    };
    dfs(0, -1);
    return ans;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
impl Solution {
    pub fn max_k_divisible_components(n: i32, edges: Vec<Vec<i32>>, values: Vec<i32>, k: i32) -> i32 {
        let n = n as usize;
        let mut g = vec![vec![]; n];
        for e in edges {
            let a = e[0] as usize;
            let b = e[1] as usize;
            g[a].push(b);
            g[b].push(a);
        }

        let mut ans = 0;

        fn dfs(i: usize, fa: i32, g: &Vec<Vec<usize>>, values: &Vec<i32>, k: i32, ans: &mut i32) -> i64 {
            let mut s = values[i] as i64;
            for &j in &g[i] {
                if j as i32 != fa {
                    s += dfs(j, i as i32, g, values, k, ans);
                }
            }
            if s % k as i64 == 0 {
                *ans += 1;
            }
            s
        }

        dfs(0, -1, &g, &values, k, &mut ans);
        ans
    }
}

Solution 2: Explicit Stack

Thinking

The whole tree sum is divisible by \(k\), so deleting a subtree whose sum is also divisible by \(k\) leaves valid components. With \(n \le 3 \times 10^4\), one bottom-up accumulation of subtree sums is enough.

Recursing from the root into each child reaches depth \(n\) on a chain and overflows the call stack. A subtree sum does not depend on the order of its children.

An explicit stack of \((node, parent, state)\) therefore runs that postorder. On entry we push the exit marker and then the children; on exit we add each child's subtree sum to this node's value, and count a component when the sum is \(0\) modulo \(k\).

We note that the problem guarantees the sum of all node values in the entire tree is divisible by \(k\). Therefore, if we remove a subtree whose sum of elements is divisible by \(k\), the sum of node values in each of the remaining connected components must also be divisible by \(k\).

An explicit stack walks the tree in postorder from the root. Each frame is \((node, parent, state)\). When a node is left, we add every child's subtree sum to this node's value. If that sum is divisible by \(k\), we increment the answer by one.

The time complexity is \(O(n)\) and the space complexity is \(O(n)\), where \(n\) is the number of nodes in the tree.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution:
    def maxKDivisibleComponents(
        self, n: int, edges: List[List[int]], values: List[int], k: int
    ) -> int:
        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        sub = [0] * n
        ans = 0
        stk = [(0, -1, 0)]
        while stk:
            i, fa, state = stk.pop()
            if state == 0:
                stk.append((i, fa, 1))
                for j in g[i]:
                    if j != fa:
                        stk.append((j, i, 0))
            else:
                s = values[i]
                for j in g[i]:
                    if j != fa:
                        s += sub[j]
                if s % k == 0:
                    ans += 1
                sub[i] = s
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class Solution {
    public int maxKDivisibleComponents(int n, int[][] edges, int[] values, int k) {
        List<Integer>[] g = new List[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int[] e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        long[] sub = new long[n];
        int ans = 0;
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, -1, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int i = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.push(new int[] {i, fa, 1});
                for (int j : g[i]) {
                    if (j != fa) {
                        stk.push(new int[] {j, i, 0});
                    }
                }
            } else {
                long s = values[i];
                for (int j : g[i]) {
                    if (j != fa) {
                        s += sub[j];
                    }
                }
                if (s % k == 0) {
                    ++ans;
                }
                sub[i] = s;
            }
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution {
public:
    int maxKDivisibleComponents(int n, vector<vector<int>>& edges, vector<int>& values, int k) {
        vector<vector<int>> g(n);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        vector<long long> sub(n);
        int ans = 0;
        vector<array<int, 3>> stk{{0, -1, 0}};
        while (!stk.empty()) {
            auto [i, fa, state] = stk.back();
            stk.pop_back();
            if (state == 0) {
                stk.push_back({i, fa, 1});
                for (int j : g[i]) {
                    if (j != fa) {
                        stk.push_back({j, i, 0});
                    }
                }
            } else {
                long long s = values[i];
                for (int j : g[i]) {
                    if (j != fa) {
                        s += sub[j];
                    }
                }
                ans += s % k == 0;
                sub[i] = s;
            }
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
func maxKDivisibleComponents(n int, edges [][]int, values []int, k int) (ans int) {
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    sub := make([]int, n)
    stk := [][3]int{{0, -1, 0}}
    for len(stk) > 0 {
        cur := stk[len(stk)-1]
        stk = stk[:len(stk)-1]
        i, fa, state := cur[0], cur[1], cur[2]
        if state == 0 {
            stk = append(stk, [3]int{i, fa, 1})
            for _, j := range g[i] {
                if j != fa {
                    stk = append(stk, [3]int{j, i, 0})
                }
            }
        } else {
            s := values[i]
            for _, j := range g[i] {
                if j != fa {
                    s += sub[j]
                }
            }
            if s%k == 0 {
                ans++
            }
            sub[i] = s
        }
    }
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
function maxKDivisibleComponents(
    n: number,
    edges: number[][],
    values: number[],
    k: number,
): number {
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const sub: number[] = Array(n).fill(0);
    let ans = 0;
    const stk: [number, number, number][] = [[0, -1, 0]];
    while (stk.length) {
        const [i, fa, state] = stk.pop()!;
        if (state === 0) {
            stk.push([i, fa, 1]);
            for (const j of g[i]) {
                if (j !== fa) {
                    stk.push([j, i, 0]);
                }
            }
        } else {
            let s = values[i];
            for (const j of g[i]) {
                if (j !== fa) {
                    s += sub[j];
                }
            }
            if (s % k === 0) {
                ++ans;
            }
            sub[i] = s;
        }
    }
    return ans;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
impl Solution {
    pub fn max_k_divisible_components(
        n: i32,
        edges: Vec<Vec<i32>>,
        values: Vec<i32>,
        k: i32,
    ) -> i32 {
        let n = n as usize;
        let mut g = vec![vec![]; n];
        for e in edges {
            let a = e[0] as usize;
            let b = e[1] as usize;
            g[a].push(b);
            g[b].push(a);
        }
        let mut sub = vec![0_i64; n];
        let mut ans = 0;
        let mut stk = vec![(0_usize, -1_i32, 0_i32)];
        while let Some((i, fa, state)) = stk.pop() {
            if state == 0 {
                stk.push((i, fa, 1));
                for &j in &g[i] {
                    if j as i32 != fa {
                        stk.push((j, i as i32, 0));
                    }
                }
            } else {
                let mut s = values[i] as i64;
                for &j in &g[i] {
                    if j as i32 != fa {
                        s += sub[j];
                    }
                }
                if s % k as i64 == 0 {
                    ans += 1;
                }
                sub[i] = s;
            }
        }
        ans
    }
}

Comments