본문 바로가기

알고리즘풀이

List를 TreeNode로 변환하기

Swift에서 [1,2,3,4,nil,nil,5] 와 같이 표현된 리스트 리터럴을 TreeNode로 변환하려면 어떻게 구현해야할까?
먼저 리스트 원소의 규칙을 정해야 한다.

  • 리스트의 순서는 트리의 노드가 순서대로 위치하는 순서다.
  • 트리 노드의 정확한 순서를 지키기 위해 노드가 존재하지 않을 경우 nil로 표현해야한다.
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
      }
}
extension TreeNode {
    static func from(_ arr: [Int?]) -> TreeNode? {
        // 예외확인, 첫번째 원소가 nil이면 nil을 그대로 반환함
        guard let firstElement = arr[0] else {
            return nil
        }
        let root = TreeNode(firstElement)
        var queue = [TreeNode]()
        queue.append(root)

        // 인자로 받은 리스트의 2번째 원소부터 시작한다.
        var index = 1
        // 너비우선탐색으로 왼쪽, 오른쪽 자식노드를 하나씩 추가한다.
        while !queue.isEmpty && index < arr.count {

            let node = queue.removeFirst()
            if let leftVal = arr[index] {
                let left = TreeNode(leftVal)
                node.left = left
                queue.append(left)
            }
            if index+1 < arr.count, let rightVal = arr[index+1]{
                let right = TreeNode(rightVal)
                node.right = right
                queue.append(right)
            }
            // 리스트를 가리키는 인덱스를 2씩 늘려준다.
            index += 2
        }

        return root
    }

}