2385.amount-of-time-for-binary-tree-to-be-infected.md

2385. Amount of Time for Binary Tree to Be Infected

Prework

We can infect the parent node from child node, but we don't have the parent node. So we can build a map to store the parent node of each node. Then we can use BFS to infect the tree.

private fun buildParent(root: TreeNode, start: Int) {
    // We find the start node at the same time
    if (root.`val` == start) startNode = root
    if (root.left != null) {
        parentMap[root.left] = root
        buildParent(root.left, start)
    }
    if (root.right != null) {
        parentMap[root.right] = root
        buildParent(root.right, start)
    }
}

DFS

private val parentMap = HashMap<TreeNode, TreeNode>()
private var startNode: TreeNode? = null

fun amountOfTime(root: TreeNode?, start: Int): Int {
    if (root == null) return 0
    buildParent(root, start)
    dfs(root, 0, HashSet())
    return totalTime
}

// Use visited set to prevent duplicate visiting.
private fun dfs(root: TreeNode?, time: Int, visited: HashSet<TreeNode>) {
    if (root == null || visited.contains(root)) return
    visited.add(root)
    totalTime = maxOf(totalTime, time)
    val newTime = time + 1
    dfs(root.left, newTime, visited)
    dfs(root.right, newTime, visited)
    dfs(parentMap[root], newTime, visited)
}

// Or equivalently, we track of the `from` node to avoid duplicate visit.
private fun dfs(root: TreeNode?, time: Int, from: TreeNode?) {
    if (root == null) return
    totalTime = maxOf(totalTime, time)
    val newTime = time + 1
    if (root.left != from) dfs(root.left, newTime, root)
    if (root.right != from) dfs(root.right, newTime, root)
    if (parentMap[root] != from) dfs(parentMap[root], newTime, root)
}

// Or equivalently, we return the time from the start node
private fun dfs(root: TreeNode?, from: TreeNode?): Int {
    if (root == null) return 0
    val left = if (root.left != null && root.left != from) dfs(root.left, root) + 1 else 0
    val right = if (root.right != null && root.right != from) dfs(root.right, root) + 1 else 0
    val parent = if (parentMap[root] != null && parentMap[root] != from) dfs(parentMap[root], root) + 1 else 0

    return maxOf(maxOf(left, right), parent)
}

// Same as above, but we use the same approach in [543. Diameter of Binary Tree](../leetcode/543.diameter-of-binary-tree.md)
private fun dfs(root: TreeNode?, source: TreeNode?): Int {
    if (root == null) return 0
    val left = if (root.left != source) dfs(root.left, root) else 0
    val right = if (root.right != source) dfs(root.right, root) else 0
    val parent = if (parentMap[root] != source) dfs(parentMap[root], root) else 0
    ans = maxOf(maxOf(left, right), parent)
    return maxOf(maxOf(left, right), parent) + 1
}

BFS

We use the same idea as the DFS, but we use BFS to traverse the tree.

private fun bfs(): Int {
    var minutes = 0 
    val queue = ArrayDeque<TreeNode>()
    val visited = HashSet<TreeNode>()
    if (startNode == null) return 0
    queue.addLast(startNode!!)
    visited.add(startNode!!)

    while (queue.isNotEmpty()) {
        val size = queue.size
        repeat(size) {
            val node = queue.removeFirst()
            if (node.left != null && node.left !in visited) {
                queue.addLast(node.left)
                visited.add(node.left)
            }
            if (node.right != null && node.right !in visited) {
                queue.addLast(node.right)
                visited.add(node.right)
            }
            if (parentMap[node] != null) {
                val p = parentMap[node]!!
                if (!visited.contains(p)) {
                    queue.addLast(p)
                    visited.add(p)
                }
            }
        }
        minutes++
    }
    // A - B - C __, we will count the last segment (redundant)
    //  +1   +1  +1
    //           (X)
    return minutes - 1
}

• Time Complexity: O(n).
• Space Complexity: O(n).