[C/C++] 백준 #2014 소수의 곱(우선 순위 큐)

우선 순위 큐와 힙(heap) 자료구조는 탐욕 알고리즘(greedy algorithm)에서 많이 쓰이는 자료구조입니다.  우선 순위 큐가 힙 자료구조를 이용해서 작성된다는 것은 이미 잘 알고 있겠죠.  그래서 삽입, 최소값(또는 최대값) 삭제가 모두 \(O(\log N)\) 시간복잡도를 가지고 있습니다.

 

priority queue

 

이번 문제는 주어진 소인수로 이루어진 수들중 K번째로 작은 수를 고르는 문제입니다.  여러가지 방식을 생각해보았는데, 저는 우선 순위 큐를 이용해서 풀었습니다.

 

1) 주어진 소수를 우선 순위 큐에 넣습니다.

2) K번째 수가 나올 때까지

  2-1) 우선 순위 큐에서 수를 빼냅니다.

  2-2) 빼낸 수에 대해서 주어진 소수를 곱해서 우선 순위 큐에 넣습니다.

 

조금이라도 효율을 좋게 하기 위해서 우선 순위 큐의 저장 개수가 K개를 넘는 경우, 그리고 들어갈 수가 가장 큰 수보다 큰 경우에는 곱해진 수를 추가하지 않습니다.

 

여전히 비효율성은 남기는 하지만, 일단 이런 방식으로 문제를 해결했습니다.

 

제가 작성한 소스입니다.

//--------------------------------------
//    baekjoon #2014 Multiplication of Prime Numbers.
//        - by Aubrey Choi
//        - created at 2019-07-19
//--------------------------------------
#include <stdio.h>
#include <queue>

int main()
{
    int n, k, p[100], max = 0, i, j;
    std::priority_queue<int,std::vector<int>,std::greater<int> > q;
    scanf("%d%d", &n, &k);
    for(i = 0; i < n; i++) { scanf("%d", p+i); q.push(p[i]); }
    max = p[n-1];
    while(--k)
    {
        long long s = q.top(); q.pop();
        for(j = 0; j < n; j++)
        {
            long long t = s*p[j];
            if(t >> 31) break;
            if(t > max) { if(q.size() > k) break; max = t; }
            q.push(t);
            if(s%p[j]==0) break;
        }
    }
    printf("%d\n", q.top());
}