[Algorithm] 백준 - 2250번 트리의 높이와 너비 (Swift)

문제 소개

트리의 가장 넓은 너비를 구하는 문제이다.

 

문제 풀이

다음과 같이 문제를 풀었다. 우선 left - root - right 순으로 순회하면 그 순서가 너비의 위치(index)가 된다. 따라서 해당 깊이에서 최소, 최대의 index를 업데이트 시켜준뒤 그 차이 +1 이 너비가 된다.
 
1. 트리 생성
2. 중위 순회
3. 루트 노드를 찾아 중위 순회
4. 결과값 구하기
 

import Foundation

let n = Int(readLine()!)!
var tree = [Int: (root: Int, node: Int, left: Int, right: Int)]()

for i in 1...n {
    tree[i] = (-1, -1, -1, -1)
}

// 트리 생성
for _ in 0..<n {
    let info = readLine()!.split(separator: " ").map { Int(String($0))! }
    let (node, left, right) = (info[0], info[1], info[2])
    
    tree[node]?.node = node
    tree[node]?.left = left
    tree[node]?.right = right
    
    tree[left]?.root = node
    tree[right]?.root = node
}

var answer = [Int: (min: Int, max: Int)]()
var idx = 1
var depthValue = 1

// 중위 순회
func inorder(_ node: Int, _ depth: Int) {
    guard let node = tree[node] else { return }
    depthValue = max(depth, depthValue)
    
    if node.left != -1 {
        inorder(node.left, depth+1)
    }
    
    if answer[depth] == nil {
        answer[depth] = (n, 1)
    }
    
    if let minValue = answer[depth]?.min, let maxValue = answer[depth]?.max {
        answer[depth]?.min = min(minValue, idx)
        answer[depth]?.max = max(maxValue, idx)
    }
    idx += 1
    
    if node.right != -1 {
        inorder(node.right, depth+1)
    }
}

var root = 0

// 루트 노드 찾아 중위 순회
for i in 1...n {
    if tree[i]?.root == -1 {
        root = i
    }
}

inorder(root, 1)

// 결과값 구하기
var depth = 0
var widthValue = 0

for i in 1...depthValue {
    if let answer = answer[i] {
        let width = answer.max - answer.min + 1
        if width > widthValue {
            depth = i
            widthValue = width
        }
    }
}

print(depth, widthValue)