코딩테스트 연습 - k진수에서 소수 개수 구하기

문제 풀이 아이디어

k진수로 변경하기

다른 언어의 경우 아래처럼 n을 k로 나눈 나머지를 가지고 구하는 경우가 많습니다.

func toRadix(_ n: Int, _ k: Int) -> String {
    var s = ""
    var n = n
    
    while n > 0 {
        s += String(n % k)
        n = n / k
    }
    return String(s.reversed())
}

하지만 Swift의 경우에는 Int를 k진수로 바꾸어서 String으로 리턴해주는 String initializer가 있습니다. 훨신 편리합니다.

String(n, radix: k)

소수 구하기

처음에는 그냥 간단하게 아래처럼 소수의 정의에 충실한(?) extension를 구현했었는데요. 이 코드를 제출하면 시간 초과가 납니다. 시간복잡도가 O(N)이기 때문입니다. n이 최대 1,000,000이지만 k 진수로 변경했을 때 얼마나 큰 수가 나올지는 예측이 불가능하기에 최대한 빠른 알고리즘을 사용하는 것이 좋습니다.

extension Int {
    var isPrime: Bool {
        guard self > 1 else { return false }
        for i in 2..<self {
            if self % i == 0 { return false }
        }
        return true
    }
}

에라토스테네스의 체

에라토스테네스의 체의 원리를 사용해서 소수를 구하는 extension입니다. n이 2로 나누어 떨어지지 않는다면 2의 배수로는 나누어 떨어지지 않습니다. 그리고 n의 최대 약수는 n의 제곱근이므로 그 이상의 수에 대해서는 나누어 떨어지는지 확인할 필요가 없습니다. 이 두 가지 원리를 통해서 O(sqrt(N)/2)까지 시간복잡도를 줄일 수 있습니다.

import Foundation

extension Int {
    var isPrime: Bool {
        // 2보다 작은 수 (1혹은 자연수가 아닌 수)는 소수가 아니다
        if self < 2 { return false }
        // 2는 소수이다.
        if self == 2 { return true }
        // 짝수는 소수가 아니다.
        if self % 2 == 0 { return false }
        
        // self의 최대 약수는 self의 제곱근 (그 보다 큰 수는 나누지지 않으므로 확인할 필요가 없음)
            // 3에서 시작해서 홀수만 나누어 떨어지는 것이 없는지 확인하면 됨 (짝수는 이미 2로 나누어봐서 확인함)
            // 시간복잡도 O(sqrt(N)/2)
        return !stride(from: 3, through: Int(sqrt(Double(self))), by: 2).contains { self % $0 == 0 }
    }
}

코드

import Foundation

extension Int {
    var isPrime: Bool {
        // 2보다 작은 수 (1혹은 자연수가 아닌 수)는 소수가 아니다
        if self < 2 { return false }
        // 2는 소수이다.
        if self == 2 { return true }
        // 짝수는 소수가 아니다.
        if self % 2 == 0 { return false }
        
        // self의 최대 약수는 self의 제곱근 (그 보다 큰 수는 나누지지 않으므로 확인할 필요가 없음)
            // 3에서 시작해서 홀수만 나누어 떨어지는 것이 없는지 확인하면 됨 (짝수는 이미 2로 나누어봐서 확인함)
            // 시간복잡도 O(sqrt(N)/2)
        return !stride(from: 3, through: Int(sqrt(Double(self))), by: 2).contains { self % $0 == 0 }
    }
}

func toRadix(_ n: Int, _ k: Int) -> String {
    var s = ""
    var n = n
    
    while n > 0 {
        s += String(n % k)
        n = n / k
    }
    return String(s.reversed())
}

func solution(_ n:Int, _ k:Int) -> Int {
    
    let arr = String(n, radix: k).split(separator: "0").map { Int($0)! }
    var cnt = 0

    for num in arr {
        if num.isPrime { cnt += 1 }
    }
    
    return cnt
}