Leetcode: Maximum Depth of Binary Tree, 이진 트리의 최대 깊이

문제설명

• 이진 트리가 주어졌을 때, 이진 트리의 루트 노드에서 리프 노드까지 가장 긴 경로의 길이를 구하라

문제풀이

• stack과 Int 리스트를 사용한다.
• stack에 노드를 담는다. stack에서 노드를 꺼내고 자식 노드가 존재한다면, 다시 stack에 담는다.
• stack에 자식 노드를 추가할 때, depth를 저장한 리스트(depths)에 해당 노드의 깊이를 저장한다.
• stack과 depths의 아이템을 동시에 꺼내게 되는데, 이때 최대 depth를 반복하여 구한다.

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
  }
}

func maxDepth(_ root: TreeNode?) -> Int {
    guard let rootNode = root else {
        return 0
    }
    var stack = [TreeNode]()
    var depths = [Int]()
    var depth = 0

    stack.append(rootNode)
    depths.append(1)

    while !stack.isEmpty {

        if let lastNode = stack.popLast(), let currentDepth = depths.popLast() {
            depth = max(depth, currentDepth)
            if let left = lastNode.left {
                stack.append(left)
                depths.append(currentDepth+1)
            }
            if let right = lastNode.right {
                stack.append(right)
                depths.append(currentDepth+1)
            }
        }
    }
    return depth
}

var depth = maxDepth(TreeNode(1,
                  TreeNode(2,
                           TreeNode(4, TreeNode(5), nil), nil),
                  TreeNode(3))) // 3
print(depth) // 3
print(maxDepth(nil)) // 0