[프로그래머스/Day 15] 소수 만들기 - 에라토스테네스의 체와 DFS를 이용한 조합 구현

오늘은 프로그래머스 '소수 만들기' 문제를 풀며 적용했던 에라토스테네스의 체 최적화 방법과, 고정된 3중 반복문을 DFS(재귀)로 확장성 있게 변경하는 과정을 정리해 보았습니다.

1. 문제 개요

이 문제는 주어진 숫자 배열 nums에 있는 숫자들 중 서로 다른 3개를 골라 더했을 때, 그 합이 소수가 되는 경우의 수를 구해야 합니다. 단순히 소수의 종류를 구하는 것이 아니라, 합이 소수가 되는 조합의 개수를 리턴해야 합니다.

---

2. 첫 번째 접근: 에라토스테네스의 체

소수를 판별할 때마다 반복문을 돌며 나누어 떨어지는지 확인하는 방식은 비효율적입니다.

주어진 숫자의 최댓값이 1,000이므로, 3개의 합은 최대 3,000을 넘지 않습니다.

따라서 에라토스테네스의 체를 사용해 소수 여부를 미리 저장해두고 O(1)로 조회하는 방식을 택했습니다.

int solution(vector<int> nums)
{
    int answer = 0;

    // 1. 에라토스테네스의 체로 소수 판별 배열 미리 만든다.
    // - nums의 원소는 최대 1000이므로, 3개의 합은 최대 3000을 넘지 않는다.
    const int MAX_SUM = 3000;

    // - 처음엔 0과 1을 제외한 모든 수가 소수라고 가정한다.
    vector<bool> isPrime(MAX_SUM + 1, true);
    isPrime[0] = false;
    isPrime[1] = false;

    // - 2부터 제곱근까지 반복하며 배수들을 제거한다.
    for(int i = 2; i * i <= MAX_SUM; ++i)
    {
        // - 특정 수가 소수라면 그 수의 배수는 소수가 아니다.
        if(isPrime[i])
        {
            for(int j = i * i; j < MAX_SUM; j += i)
            {
                isPrime[j] = false;
            }
        }
    }

    // 2. 3중 반복문으로 3개의 수를 더한다.
    const size_t NUMS_SIZE = nums.size();

    for(size_t t1 = 0; t1 < NUMS_SIZE; t1++)
    {
        for(size_t t2 = t1 + 1; t2 < NUMS_SIZE; t2++)
        {
            for(size_t t3 = t2 + 1; t3 < NUMS_SIZE; t3++)
            {
                int sum = nums[t1] + nums[t2] + nums[t3];

                if (isPrime[sum]) {
                    answer++;
                }
            }
        }
    }

    return answer;
}

 

기능적으로는 정답이지만, solution 함수 안에 소수 체 생성 로직과 조합 탐색 로직이 섞여 있어 코드가 길어지고 가독성이 떨어집니다. 만약 소수 판별 로직이 변경된다면 메인 함수를 건드려야 하는 문제가 있습니다.

---

3. 두 번째 접근: 함수 분리

"한 함수는 하나의 기능만 한다"는 원칙에 따라, 소수 판별 배열을 생성하는 로직을 별도 함수로 분리하여 코드의 응집도를 높였습니다.

// 1. 에라토스테네스의 체로 소수 판별 배열 미리 만든다.
vector<bool> getPrimeSieve(int maxNum)
{
    // - 처음엔 0과 1을 제외한 모든 수가 소수라고 가정한다.
    vector<bool> isPrime(maxNum + 1, true);
    isPrime[0] = false;
    isPrime[1] = false;

    // - 2부터 제곱근까지 반복하며 배수들을 제거한다.
    for(int i = 2; i * i <= maxNum; ++i)
    {
        // - 특정 수가 소수라면 그 수의 배수는 소수가 아니다.
        if(isPrime[i])
        {
            for(int j = i * i; j < maxNum; j += i)
            {
                isPrime[j] = false;
            }
        }
    }

    return isPrime;
}


int solution(vector<int> nums)
{
    int answer = 0;

    // 1. 에라토스테네스의 체로 소수 판별 배열 미리 만든다.
    // - nums의 원소는 최대 1000이므로, 3개의 합은 최대 3000을 넘지 않는다.
    const int MAX_SUM = 3000;

    vector<bool> isPrime = getPrimeSieve(MAX_SUM);

    // 2. 3중 반복문으로 3개의 수를 더한다.
    const size_t NUMS_SIZE = nums.size();

    for(size_t t1 = 0; t1 < NUMS_SIZE; t1++)
    {
        for(size_t t2 = t1 + 1; t2 < NUMS_SIZE; t2++)
        {
            for(size_t t3 = t2 + 1; t3 < NUMS_SIZE; t3++)
            {
                int sum = nums[t1] + nums[t2] + nums[t3];

                if (isPrime[sum]) {
                    answer++;
                }
            }
        }
    }

    return answer;
}

---

4. 세 번째 접근: 확장성을 고려한 DFS 풀이

만약 문제 조건이 "3개의 수"가 아니라 "5개 혹은 K개의 수"를 더하는 것으로 바뀐다면 어떻게 될까요?

기존의 3중 for문 방식은 for문을 K개만큼 중첩해야 하므로 유연한 대처가 불가능합니다.

이때 DFS(재귀/백트래킹)를 사용하면 확장성 있는 코드를 작성할 수 있습니다.

// [기능 1] 소수 판별 체 (이전과 동일)
vector<bool> getPrimeSieve(int maxNum)
{
    vector<bool> isPrime(maxNum + 1, true);
    isPrime[0] = false;
    isPrime[1] = false;
    
    for (int i = 2; i * i <= maxNum; i++)
    {
        if (isPrime[i])
        {
            for (int j = i * i; j <= maxNum; j += i) isPrime[j] = false;
        }
    }
    
    return isPrime;
}

// [기능 2] DFS를 이용한 조합 탐색
// idx: 현재 탐색할 nums의 인덱스 (중복 방지용)
// count: 현재까지 뽑은 숫자의 개수
// currentSum: 현재까지 뽑은 숫자들의 합
// nums, isPrime, answer: 참조(&)로 전달하여 메모리 절약 및 값 변경 공유
void dfs(int idx, int count, int currentSum, 
         const vector<int>& nums, const vector<bool>& isPrime, int& answer)
{

    // 1. 종료 조건: 3개를 모두 뽑았을 때
    if (count == 3)
    {
        if (isPrime[currentSum])
        {
            answer++;
        }
        
        return; // 더 이상 깊이 들어가지 않고 돌아갑니다.
    }

    // 2. 재귀 로직: 현재 인덱스(idx)부터 끝까지 순회하며 하나를 선택
    for (int i = idx; i < nums.size(); i++)
    {
        // nums[i]를 선택하고, 다음 재귀에서는 i+1부터 탐색하도록 넘김
        dfs(i + 1, count + 1, currentSum + nums[i], nums, isPrime, answer);
    }
}

int solution(vector<int> nums)
{
    int answer = 0;
    const int MAX_SUM = 3000;

    // 소수 체 준비
    vector<bool> isPrime = getPrimeSieve(MAX_SUM);

    // DFS 탐색 시작 (0번 인덱스부터, 0개 뽑음, 합계 0)
    dfs(0, 0, 0, nums, isPrime, answer);

    return answer;
}

---

5. 정리

1. 알고리즘 최적화: 반복적인 소수 판별이 필요할 땐 에라토스테네스의 체를 사용하여 시간 복잡도를 O(1)로 줄인다.
2. 기능 분리: 핵심 알고리즘(체 생성)과 비즈니스 로직(문제 풀이)을 분리하여 코드 가독성과 재사용성을 높인다.
3. 확장성 확보: 고정된 N중 반복문 대신 DFS(백트래킹)를 사용하면, 선택해야 할 원소의 개수가 변하더라도 유연하게 대응할 수 있다.