Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 K번째로 작은 소수 분수 찾기

문제 개요

정렬된 리스트가 하나 주어진다고 가정해 봅시다. 이 리스트에는 1과 여러 개의 소수가 포함되어 있습니다. 리스트 내의 모든 p < q 조합에 대해 분수 p/q를 고려할 때, 이 분수들 중 K번째로 작은 분수를 찾아야 합니다.

결과는 배열 형태로 반환하며, ans[0]에는 분자 p를, ans[1]에는 분모 q를 담습니다.

예를 들어 입력이 [1, 3, 5, 7]이고 k = 2라고 해봅시다. 만들 수 있는 분수는 1/3, 1/5, 1/7, 3/5, 3/7, 5/7이고, 이 중 두 번째로 작은 값은 1/5입니다. 따라서 정답은 [1, 5]가 됩니다.

풀이 접근 방법

이 문제는 최소 힙(Min-Heap) 역할을 하는 우선순위 큐(priority queue)를 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 각 분모에 대해 가장 작은 분수부터 차례대로 후보군에 넣고, K번 pop할 때까지 다음 후보를 계속 추가하는 것입니다.

구체적인 단계는 다음과 같습니다.

  • 분자 a, 분모 b, 그리고 a/b 값을 저장하는 Data 구조체를 정의합니다.
  • 크기가 2인 정답 배열 ret을 선언합니다.
  • n := 배열 A의 크기로 설정합니다.
  • Data 타입을 저장하는 우선순위 큐 pq를 생성합니다.
  • i := 0부터 n 미만까지 반복하면서 Data(A[0], A[i], 0)를 pq에 삽입합니다. 즉, 모든 분모에 대해 분자가 A[0](=1)인 분수를 먼저 넣습니다.
  • K가 0이 될 때까지 아래 과정을 반복합니다.
    • pq의 최상단 원소를 temp로 꺼내고(pop) 큐에서 제거합니다.
    • K가 0이 되었다면, ret[0]에 temp의 a(분자), ret[1]에 temp의 b(분모)를 저장하고 ret을 반환합니다.
    • temp.idx + 1 < n이라면, idx := temp.idx + 1로 갱신한 뒤 Data(A[idx], temp.b, idx)를 pq에 삽입합니다. 같은 분모에 대해 더 큰 분자를 가진 다음 후보를 추가하는 것입니다.
  • 반복이 끝나면 ret을 반환합니다.

C++ 구현 예제

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
   cout << "[";
   for(int i = 0; i<v.size(); i++){
      cout << v[i] << ", ";
   }
   cout << "]"<<endl;
}
struct Data{
   double val, a, b;
   int idx;
   Data(double a, double b, int c){
      val = a / b;
      this->a = a;
      this->b = b;
      idx = c;
   }
};
struct Comparator{
   bool operator()(Data a, Data b){
      return !(a.val < b.val);
   }
};
class Solution {
public:
   vector<int> kthSmallestPrimeFraction(vector<int>& A, int K) {
      vector <int> ret(2);
      int n = A.size();
      priority_queue <Data, vector <Data>, Comparator> pq;
      for(int i = 0; i < n; i++){
         pq.push(Data(double(A[0]), double(A[i]), 0));
      }
      while(K--){
         Data temp = pq.top();
         pq.pop();
         if(K == 0){
            ret[0] = temp.a;
            ret[1] = temp.b;
            return ret;
         }
         if(temp.idx + 1 < n){
            int idx = temp.idx + 1;
            pq.push(Data(double(A[idx]), double(temp.b), idx));
         }
      }
      return ret;
   }
};
main(){
   Solution ob;
   vector<int> v = {1,3,5,7};
   print_vector(ob.kthSmallestPrimeFraction(v, 2));
}

입력

{1,3,5,7}
2

출력

[1, 5]

동작 원리 설명

처음에는 분자를 1로 고정하고 각 소수를 분모로 하는 분수들을 우선순위 큐에 넣습니다. 이후 매 반복마다 현재 가장 작은 분수를 꺼내고, 해당 분모에 대해 다음으로 큰 분자를 가진 분수를 새로 삽입합니다. 이렇게 하면 전체 분수를 모두 생성하지 않고도 필요한 만큼만 작은 분수부터 순서대로 확인할 수 있어 시간 복잡도를 O(K log N) 수준으로 줄일 수 있습니다.

Comparator에서 !(a.val < b.val)를 반환하는 이유는 C++의 priority_queue가 기본적으로 최대 힙으로 동작하기 때문입니다. 비교 연산을 뒤집어 실질적인 최소 힙처럼 사용하는 것입니다.