백준 3078번 좋은 친구 스위프트(Swift)

가끔씩 독자분들에게 조언과 의견을 구하기 위해서 해결하기 힘들었던 코딩 테스트 문제 풀이를 공유해볼까합니다.

 

https://www.acmicpc.net/problem/3078

 

큐를 이용하여 푸는 문제였다.

테스트케이스도 통과하고 답도 맞으나, 사이트에서는 시간 초과.

처음의 O(n^2)의 시간복잡도를 가진 코드에서 다양하게 개선시켜보았으나 여전히 시간초과였다.

고질적인 백준의 Swift 한정 readLine() 시간초과가 원인이 아닐까 싶다. 다른 사이트들은 swift는 꽤 널널하게 잡아주거나 하던데... 참 이게 기업코테에서도 문제가 된다면 어떻게 해야할지가 요새 제일 큰 고민이다.

 

다음은 O(n^2)의 시간복잡도를 가진 풀이이다.

struct Queue<T> {
    private var enqueue: [T?]
    private var dequeue: [T?] = []
    
    var isEmpty: Bool{
        return enqueue.isEmpty && dequeue.isEmpty
    }
    var count: Int{
        return enqueue.count + dequeue.count
    }
    
    init() {
        enqueue = [T?]()
    }
    
    mutating func push(_ element: T){
        enqueue.append(element)
    }
    
    mutating func pop() -> T?{
        if dequeue.isEmpty{
            dequeue = enqueue.reversed()
            enqueue.removeAll()
        }
        return dequeue.popLast()!
    }
    
}

func countGoodFriends(_ students: Queue<(String, Int)>, _ K: Int) -> Int {
    var goodFriendsCount = 0
    
    var queue = students
    var queueLeft = Queue<(String, Int)>()
    
    for _ in 0..<students.count {
        let firstStudent = queue.pop()
        queueLeft = queue
        
        
        while !queueLeft.isEmpty {
            if let comparedStudent = queueLeft.pop() {
                if firstStudent!.0.count == comparedStudent.0.count && abs(firstStudent!.1 - comparedStudent.1) <= K {
                    goodFriendsCount += 1
                }
            }
        }
        
    }
    
    return goodFriendsCount
}

let input = readLine()!.split(separator: " ").map { Int($0)! }
let N = input[0]
let K = input[1]

var students = Queue<(String, Int)>()
for i in 1...N {
    let name = readLine()!
    students.push((name,i))
}

let result = countGoodFriends(students, K)
print(result)

 

다음은 조금 더 개선된 코드입니다만, 여전히 시간초과입니다.

 

func countGoodFriends(_ students: Queue<(String, Int)>, _ K: Int) -> Int {
    var studentArr = students
    var goodFriendsCount = 0
    var nameLengthToGrades = [Int: [Int]]()

    for _ in 0..<students.count {
        let student = studentArr.pop()

        if let grades = nameLengthToGrades[student!.0.count] {
            for grade in grades {
                if abs(student!.1 - grade) <= K {
                    goodFriendsCount += 1
                }
            }
        }
        
        nameLengthToGrades[student!.0.count, default: []].append(student!.1)
    }

    return goodFriendsCount
}

let input = readLine()!.split(separator: " ").map { Int($0)! }
let N = input[0]
let K = input[1]

var students = Queue<(String, Int)>()
for i in 1...N {
    let name = readLine()!
    students.push((name,i))
}

let result = countGoodFriends(students, K)
print(result)

 

그리고 발견한 다음의 글...
https://velog.io/@rarebook92/%EB%B0%B1%EC%A4%80%EC%9D%80-Swift%EC%97%90%EA%B2%8C-%EA%B4%80%EB%8C%80%ED%95%B4%EB%9D%BC

이 풀이에 대한 다양한 의견과 조언을 받고 있습니다.