문제설명
- 이진트리가 주어졌을 때, 트리를 level 순서로 순회한 노드의 값들을 반환하라
- 순서는 왼쪽에서 오른쪽, 위에서 아래로가 된다.
문제풀이
- 이진 트리를 너비우선탐색(BFS)로 순회한다.
- 각 레벨마다 리스트에 리스트를 추가한다.
- 각 레벨은 어떻게 알 수 있을까? 큐에 남아있는 갯수만큼 다돌면 하나의 레벨을 순회한 것이다.
// MARK: - Answer
func levelOrder(_ root: TreeNode?) -> [[Int]] {
var result = [[Int]]()
guard let root = root else {
return result
}
var queue = [TreeNode]()
queue.append(root)
var level = 0
while !queue.isEmpty {
let queueLength = queue.count
result.append(Array<Int>())
for _ in 0..<queueLength {
let node = queue.removeFirst()
result[level].append(node.val)
if let left = node.left {
queue.append(left)
}
if let right = node.right {
queue.append(right)
}
}
level+=1
}
return result
}
// MARK: - Test
var node1: TreeNode?
node1 = TreeNode.from([1,2,2,3,4,4,3])
print(levelOrder(node1)) // [[1],[2,2],[3,4,4,3]]
node1 = TreeNode.from([1,2,2,nil,3,nil,3])
print(levelOrder(node1)) // [[1],[2,2],[3,3]]
// MARK: - Definition
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? {
guard let firstElement = arr[0] else {
return nil
}
let root = TreeNode(firstElement)
var queue = [TreeNode]()
queue.append(root)
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)
}
index += 2
}
return root
}
}