LeetCode in Kotlin

2385. Amount of Time for Binary Tree to Be Infected

Medium

You are given the root of a binary tree with unique values, and an integer start. At minute 0, an infection starts from the node with value start.

Each minute, a node becomes infected if:

Return the number of minutes needed for the entire tree to be infected.

Example 1:

Input: root = [1,5,3,null,4,10,6,9,2], start = 3

Output: 4

Explanation: The following nodes are infected during:

It takes 4 minutes for the whole tree to be infected so we return 4.

Example 2:

Input: root = [1], start = 1

Output: 0

Explanation: At minute 0, the only node in the tree is infected so we return 0.

Constraints:

Solution

import com_github_leetcode.TreeNode

/*
 * Example:
 * var ti = TreeNode(5)
 * var v = ti.`val`
 * Definition for a binary tree node.
 * class TreeNode(var `val`: Int) {
 *     var left: TreeNode? = null
 *     var right: TreeNode? = null
 * }
 */
class Solution {
    private var max = 0
    fun amountOfTime(root: TreeNode?, start: Int): Int {
        dfs(root, start, Distance(-1))
        return max
    }

    private fun dfs(root: TreeNode?, start: Int, l: Distance): Int {
        if (root == null) {
            return 0
        }
        val ld = Distance(-1)
        val rd = Distance(-1)
        val left = dfs(root.left, start, ld)
        val right = dfs(root.right, start, rd)
        if (l.`val` == -1 && start == root.`val`) {
            max = Math.max(left, right)
            l.`val` = 1
        }
        if (ld.`val` != -1) {
            max = Math.max(max, ld.`val` + right)
            l.`val` = ld.`val` + 1
        } else if (rd.`val` != -1) {
            max = Math.max(max, rd.`val` + left)
            l.`val` = rd.`val` + 1
        }
        return Math.max(left, right) + 1
    }

    private class Distance internal constructor(var `val`: Int)
}