2872. Maximum Number of K-Divisible Components

Depth-first search (DFS) - depth-first search is always a difficult algorithm to visualize and understand the solution because when traversing from 0 to n, the results will be from n to 0, and this is one example.

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 = 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.

Solution Approach:

- Trim leaves or branches where the total from the root of that branch and its child leaves/branches is divisible by k.

Coding:

- Use depth-first search (DFS), starting from any node from 0 to n-1.

- Mark a leaf with a value that is divisible by k or a branch where the sum of the root and its child leaves/branches is divisible by k.

- Count the marked points to get the desired result.

func maxKDivisibleComponents(n int, edges [][]int, values []int, k int) int {
	//store tree by mapEdges
	mapEdges := make(map[int]map[int]bool)
	for i := range edges {
		if mapEdges[edges[i][0]] == nil {
			mapEdges[edges[i][0]] = make(map[int]bool)
		}
		if mapEdges[edges[i][1]] == nil {
			mapEdges[edges[i][1]] = make(map[int]bool)
		}
		mapEdges[edges[i][0]][edges[i][1]] = true
		mapEdges[edges[i][1]][edges[i][0]] = true
	}
	// mapCheck check which branch was processed
	mapCheck := make(map[int]bool)
	// can go with any values form 0 to n-1
	_, rs := calculateBranch(mapEdges, mapCheck, n-1, values, k)
	return rs
}
func calculateBranch(mapEdges map[int]map[int]bool, mapCheck map[int]bool, r int, values []int, k int) (int, int) {
	sum := values[r]
	mark := 0
	mapCheck[r] = true
	for ri := range mapEdges[r] {
		if !mapCheck[ri] {
			//calculate sub branch
			sumi, marki := calculateBranch(mapEdges, mapCheck, ri, values, k)
			mark += marki
			sum += sumi
		}
	}
	//mark if branch r can be cut or not
	if sum%k == 0 {
		mark++
	}
	return sum, mark
}