Algorithm / 소수 만들기

1) 문제

주어진 숫자 중 3개의 수를 더했을 때 소수가 되는 경우의 개수를 구하려고 합니다. 숫자들이 들어있는 배열 nums가 매개변수로 주어질 때, nums에 있는 숫자들 중 서로 다른 3개를 골라 더했을 때 소수가 되는 경우의 개수를 return 하도록 solution 함수를 완성해주세요.

제한사항

• nums에 들어있는 숫자의 개수는 3개 이상 50개 이하입니다.
• nums의 각 원소는 1 이상 1,000 이하의 자연수이며, 중복된 숫자가 들어있지 않습니다.

2) 문제 풀이

나의 말:
func solution(_ nums:[Int]) -> Int {
    var answer: Int = 0
    var index: Int
    var dex: Int
    
    for (i, num) in nums.enumerated() {
        index = i + 1 // 초기화
        dex = index + 1
        
        while index + 1 < nums.count {
            
            let result = num + nums[index] + nums[dex]
                        
            var isPrime = true
            
            for i in stride(from: 2, to: result, by: 1) where result % i == 0 {
                isPrime = false
                break
            }
            
            if isPrime {
                answer += 1
            }
            
            if dex == nums.count - 1 {
                index += 1
                dex = index + 1
            } else {
                dex += 1
            }
            
        }
    }

    return answer
}

결과

3) 코드 개선

• 중복 조합 처리 미흡 문제 → 모든 숫자 조합을 정확하게 순회하지 못할 가능성이 있음
• 가독성 문제 → while 내부 흐름이 직관적이지 않아 가독성이 떨어짐
• 구조를 DFS 방식으로 변환 → 네스팅이 낮아지고 가독성 향상

func solution(_ nums: [Int]) -> Int {
    var count = 0

    func dfs(_ index: Int, _ depth: Int, _ sum: Int) {
        if depth == 3 {
            if isPrime(sum) {
                count += 1
            }
            return
        }

        for i in index..<nums.count {
            dfs(i + 1, depth + 1, sum + nums[i])
        }
    }

    dfs(0, 0, 0)
    return count
}

func isPrime(_ num: Int) -> Bool {
    if num < 2 { return false }
    for i in 2...Int(Double(num).squareRoot()) {
        if num % i == 0 {
            return false
        }
    }
    return true
}

결과