BinaryTree, PrefixSum_Leetcode_2385. Amount of Time for Binary Tree to Be Infected [swift]

https://leetcode.com/problems/amount-of-time-for-binary-tree-to-be-infected/description/

 

문제 접근

아이디어는 두 가지가 생각났다.

1. node.val은 유니크하다. 그러므로 다익스트라 알고리즘으로 풀기
2. 최하단 노드로부터 depth를 계산하기 (dfs)

정답 코드

양방향 노드, 우선순위 Queue로 인해서 O(nlogn)

처참한 성적

642msBeats5.55% -> 95/100 ;;; 95등이라니!

54.86MBBeats27.78%

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     public var val: Int
 *     public var left: TreeNode?
 *     public var right: TreeNode?
 *     public init() { self.val = 0; self.left = nil; self.right = nil; }
 *     public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
 *     public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
 *         self.val = val
 *         self.left = left
 *         self.right = right
 *     }
 * }
 */

class Solution {
    var graph = [Int: [Int]]()

    func amountOfTime(_ root: TreeNode?, _ start: Int) -> Int {
        makeGraph(root)
        var distance = [Int](repeating: Int.max, count: start + 1)
        var queue = [(Int,Int)]()

        distance[start] = 0
        queue.append((start,0))

        while !queue.isEmpty {
            let (fromIdx, historyDst) = queue.removeFirst()
            for end in graph[fromIdx, default: []] {
                if end >= distance.count {
                    let repeatValue = end - distance.count + 1
                    for _ in 0..<repeatValue {
                        distance.append(Int.max)
                    }
                }
                let nextDst = historyDst + 1
                if distance[end] > nextDst {
                    distance[end] = nextDst
                    queue.append((end, nextDst))
                    queue.sort(by: {
                        $0.1 < $1.1
                    })
                }
            }

        }
        
        return distance.filter { $0 != Int.max }.max() ?? 0 
    }

    func makeGraph(_ node: TreeNode?) {
        guard let node else { return }
        if let lhs = node.left {
            graph[node.val, default: []].append(lhs.val)
            graph[lhs.val, default: []].append(node.val)
            
            makeGraph(lhs)
        }
        if let rhs = node.right {
            graph[node.val, default: []].append(rhs.val)
            graph[rhs.val, default: []].append(node.val)
            
            makeGraph(rhs)
        }
    }
}

모범 코드를 보고 공부해서 다시 작성

dfs방식이 옳았다.

최상단 root node를 기준으로 생각한다.

rootnode를 기준으로 두 자식노드 중 한 쪽에서는 start로 부터의 값이 필요하고, 다른 한 쪽에서는 전체 노드의 길이가 필요하다.

자식이 없는 노드는 최하단 노드이다.

 

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     public var val: Int
 *     public var left: TreeNode?
 *     public var right: TreeNode?
 *     public init() { self.val = 0; self.left = nil; self.right = nil; }
 *     public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
 *     public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
 *         self.val = val
 *         self.left = left
 *         self.right = right
 *     }
 * }
 */
class Solution {
    var maxValue = 0

    func amountOfTime(_ root: TreeNode?, _ start: Int) -> Int {
        dfs(root, start)
        return maxValue
    }

    func dfs(_ root: TreeNode?, _ start: Int) -> (depth: Int, fromStart: Int?) {
        guard let root else { return (0, nil) }
        
        let (lhsDepth, lhsFromStart) = dfs(root.left, start)
        let (rhsDepth, rhsFromStart) = dfs(root.right, start)

        let maxFromBoth = max(lhsDepth, rhsDepth)
        // 찾은 노드가 startnode
        if root.val == start {
            maxValue = max(maxValue, maxFromBoth)
            return (maxFromBoth + 1 ,0)
        }

        if let lhsFromStart {
            let currentDepth = rhsDepth + lhsFromStart + 1
            maxValue = max(currentDepth, maxValue)
            return (maxFromBoth + 1 ,lhsFromStart + 1)
            // return (-1,lhsFromStart + 1)
        }

        if let rhsFromStart {
            let currentDepth = lhsDepth + rhsFromStart + 1
            maxValue = max(currentDepth, maxValue)
            return (maxFromBoth + 1 ,rhsFromStart + 1)
            // return (-1 ,rhsFromStart + 1)
        }

        
        return (maxFromBoth + 1, nil)
    }
}

 

주석처리된 return 문은 궁금해서 작동실험해봤던 흔적이다.

서브트리 중, start가 존재했던 subTree의 depth는 불필요하기에 궁금해서 넣어봤다.

물론 잘 작동한다.