[BOJ, Swift] 17298 - 오큰수 ( with 스택 )

17298번: 오큰수

배열에서 자기 자신보다 오른쪽에 있는 가장 가까운 큰 수를 찾는 문제이다. 처음 구현한 방법은 이중 for문을 활용하였으나, 배열의 크기가 클 땐 O(n^2) 의 시간복잡도를 감당할 수 없어 시간초과가 나왔다.

실패 코드

import Foundation

let n = Int(readLine()!)!
var numbers = readLine()!.split { $0 == " " }.map { Int(String($0))! }
var result : [Int] = []

for i in 1 ..< n {
	// 바로 옆의 값이 크면 그게 바로 오큰수이므로 정답 배열에 삽입
    if numbers[i - 1] < numbers[i] {
        result.append(numbers[i])
    } else {
    	
        // 아니라면 그 다음 값부터 큰 값을 조사
        var idx = 0
        for j in i + 1 ..< n {
            if numbers[i - 1] < numbers[j] {
            	// 나오는 순간 정답 배열 삽입 후 Break
                result.append(numbers[j])
                break
            } else {
            	// 안나온다면 위치를 기억해둠
                idx += 1
            }
        }
        
        // 위치가 순회한 길이와 같다면 오큰수가 없다고 판단, -1 삽입
        if idx == n - (i + 1) {
            result.append(-1)
        }
    }
}

// 마지막은 무조건 -1
result.append(-1)
print(result.map { String(Int($0)) }.joined(separator: " "))

 

따라서 자료구조 스택을 활용하여 O(n) 의 시간복잡도로 해결하였다. 후입선출(LIFO)의 자료구조인 스택은 배열을 활용하였다.

AC 코드

let size = Int(readLine()!)!
let seq = readLine()!.split { $0 == " " }.map { Int(String($0))! }
var ans = Array(repeating: 0, count: size)
var stack = Array<Int>()

// 배열을 역으로 순회(idx 0 포함)
for i in stride(from: size - 1, through: 0, by: -1) {
    
    // 스택이 안비어있고, 현재 위치에서 크기가 스택의 마지막보다 크면
    while !stack.isEmpty && seq[i] >= stack.last! {
        
        // 스택의 마지막을 버려가며 가장 큰 수를 찾는다
        stack.removeLast()
    }
    
    // 그랬는데도 비어있다면, 오큰수가 없으므로 -1
    if stack.isEmpty {
        ans[i] = -1
    } else { // 있다면 스택의 마지막으로 값 갱신
        ans[i] = stack.last!
    }
    
    // sample[i]의 값이 다음 녀석에겐 오큰수가 될 수 있으므로 넣어줌
    stack.append(seq[i])
}

print(ans.map { String(Int($0)) }.joined(separator: " "))