LeetCode in Kotlin

1443. Minimum Time to Collect All Apples in a Tree

Medium

Given an undirected tree consisting of n vertices numbered from 0 to n-1, which has some apples in their vertices. You spend 1 second to walk over one edge of the tree. Return the minimum time in seconds you have to spend to collect all apples in the tree, starting at vertex 0 and coming back to this vertex.

The edges of the undirected tree are given in the array edges, where edges[i] = [ai, bi] means that exists an edge connecting the vertices ai and bi. Additionally, there is a boolean array hasApple, where hasApple[i] = true means that vertex i has an apple; otherwise, it does not have any apple.

Example 1:

Input: n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false,false,true,false,true,true,false]

Output: 8

Explanation: The figure above represents the given tree where red vertices have an apple. One optimal path to collect all apples is shown by the green arrows.

Example 2:

Input: n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false,false,true,false,false,true,false]

Output: 6

Explanation: The figure above represents the given tree where red vertices have an apple. One optimal path to collect all apples is shown by the green arrows.

Example 3:

Input: n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false,false,false,false,false,false,false]

Output: 0

Constraints:

Solution

@Suppress("UNUSED_PARAMETER")
class Solution {
    fun minTime(n: Int, edges: Array<IntArray>, hasApple: List<Boolean>): Int {
        val visited: MutableSet<Int> = HashSet()
        val graph: MutableMap<Int, MutableList<Int>> = HashMap()
        for (edge in edges) {
            val vertexA = edge[0]
            val vertexB = edge[1]
            graph.computeIfAbsent(vertexA) { _: Int? -> ArrayList() }.add(vertexB)
            graph.computeIfAbsent(vertexB) { _: Int? -> ArrayList() }.add(vertexA)
        }
        visited.add(0)
        val steps = helper(graph, hasApple, 0, visited)
        return if (steps > 0) steps - 2 else 0
    }

    private fun helper(
        graph: Map<Int, MutableList<Int>>,
        hasApple: List<Boolean>,
        node: Int,
        visited: MutableSet<Int>
    ): Int {
        var steps = 0
        for (child in graph.getOrDefault(node, mutableListOf())) {
            if (visited.contains(child)) {
                continue
            } else {
                visited.add(child)
            }
            steps += helper(graph, hasApple, child, visited)
        }
        return if (steps > 0) {
            steps + 2
        } else if (java.lang.Boolean.TRUE == hasApple[node]) {
            2
        } else {
            0
        }
    }
}