LeetCode알고리즘 Kth Smallest Element in a BST문제

오늘도 IDE 없이 푸는 연습을 해보았습니다.
문제가 조금이라도 어려워지면 디버깅이 필요해서 IDE 없이 푸는 게 어렵더라고요.
그래서 차근차근 연습하려 합니다.

오늘 풀 문제는 230번이고 중간 난이도고 62.8%의 성공률을 보이네요.
Kth Smallest Element in a BST인데 (이진 탐색 트리) BST가 BTS로 보이는 거..ㅎㅎ뭐지

문제를 살펴봅시다.

이진 탐색 트리의 루트와 k가 주어지면 트리에서 k번째로 작은 걸 리턴해주면 됩니다.

예시를 보면 이해가 확 갈 겁니다요.

제가 푸려는 방법은 중위 순회를 하면서 배열에 값을 넣고 배열에서 k번째를 리턴해주는 방법으로 풀어보려 합니다.

그리고 문제에서는 아래와 같이 노드에 대한 구현이 주어집니다.

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

먼저 루트가 닐인지 체크를 해주고 닐이면 0을 리턴해줍니다.
answer변수가 순회를 하면서 값을 저장할 배열입니다.
순회라는 함수를 두었고 여기 안에서 순회를 하게 됩니다.
중위 순회를 하기 위해서 루트의 왼쪽, 중간, 오른쪽 순으로 값을 넣어줍니다.
순회 함수를 외부로 빼고 answer를 inout하게 하는 방법도 있지만 간단하니 함수 내 함수로 구현했습니다.
그리고 마지막으로 순회된 값이 저장된 answer에서 k번째 - 1인 값을 리턴해줍니다.
(인덱스가 0부터 시작하니 k번째는 k - 1에 저장된 값이 됩니다.)

func kthSmallest(_ root: TreeNode?, _ k: Int) -> Int {
    guard let node = root else {
        return 0
    }
    
    var answer = [Int]()
                
    func traversal(_ root: TreeNode?) {
        guard let node = root else {
            return
        }

        traversal(node.left)
        answer.append(node.val)
        traversal(node.right)
    }
    
    traversal(node)
    
    return answer[k-1]
}

여기까지 해서 제출을 하면 아래와 같은 결과를 얻을 수 있습니다.

 

resource: leetcode.com/problems/kth-smallest-element-in-a-bst/