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

C++로 해결하는 정렬된 행렬의 K번째 최소 배열 합 문제

문제 개요

m × n 크기의 행렬 mat과 정수 k가 주어졌다고 가정해 봅시다. 이 행렬의 각 행은 비내림차순(오름차순)으로 정렬되어 있습니다. 우리는 각 행에서 정확히 하나의 원소를 선택하여 하나의 배열을 만들 수 있으며, 이렇게 만들 수 있는 모든 배열의 합 중에서 K번째로 작은 합을 찾아야 합니다.

예제로 이해하기

입력이 다음과 같다고 가정해 보겠습니다.

mat = [[1,3,11],[2,4,6]]
1311
246

이때 k = 5라면 출력은 7이 됩니다. 각 행에서 하나씩 원소를 선택했을 때 가장 작은 합 5개는 순서대로 [1,2], [1,4], [3,2], [3,4], [1,6]이며, 그중 5번째 합이 바로 7이기 때문입니다.

풀이 접근 방법

이 문제는 우선순위 큐(priority queue)를 활용하면 효율적으로 해결할 수 있습니다. 전체 조합을 모두 생성하지 않고, 현재까지의 최솟값부터 차례대로 후보를 확장해 나가는 방식입니다.

  • 최소 합 기준으로 동작하는 우선순위 큐 pq를 정의합니다.

  • 행렬 정보를 저장할 2차원 배열 m을 정의합니다.

  • 후보 조합을 확장하는 함수 update()를 정의합니다. 이 함수는 배열 v, 인덱스 i, 불리언 값 ok(기본값 false)를 매개변수로 받습니다.

  • i가 v의 크기와 같다면, 즉 모든 행에 대한 인덱스가 결정되었다면:

    • ok가 false이면 아무 변화 없이 종료합니다.

    • 그렇지 않으면 각 행에서 선택한 원소들의 합을 계산합니다.

    • v를 복사한 temp 배열을 만들고, 맨 앞에 합을 삽입한 뒤 pq에 넣습니다.

  • v[i]를 1 증가시킨 경우와 증가시키지 않은 경우를 모두 재귀적으로 탐색하여, 인덱스 조합의 분기를 처리합니다.

  • 메인 로직에서는 다음을 수행합니다.

    • m에 주어진 행렬을 저장하고, ret을 0으로 초기화합니다.

    • n은 행의 개수, z는 열의 개수로 설정합니다.

    • 각 행의 첫 번째 원소(각 행의 최솟값)를 ret에 더합니다. 이 값이 가장 작은 초기 합이 됩니다.

    • 크기가 n인 temp 배열을 만들어 초기 합을 맨 앞에 넣고 pq에 삽입합니다.

    • 중복 방지를 위한 집합 s를 정의합니다.

    • k번 반복하면서 매번 pq의 최상단(현재 최소 합)을 꺼내 ret에 저장하고, 해당 조합에서 인덱스를 하나씩 이동시켜 새로운 후보들을 update()로 pq에 추가합니다.

    • 이미 집합 s에 있는 중복 조합이 pq의 상단에 있다면 제거합니다.

  • 반복이 끝나면 ret을 반환합니다. 이것이 K번째로 작은 배열 합입니다.

C++ 구현 예제

아래 코드를 통해 더 명확하게 이해할 수 있습니다.

예제 코드

#include <bits/stdc++.h>
using namespace std;
struct Cmp{
   bool operator()(vector <int>& a, vector <int>& b) {
      return !(a[0] < b[0]);
   }
};
class Solution {
   public:
   priority_queue<vector<int>, vector<vector<int> >, Cmp> pq;
   vector<vector<int> > m;
   int z;
   void update(vector<int>& v, int i, bool ok = false){
      if (i == v.size()) {
         if (!ok)
         return;
         int sum = 0;
         for (int j = 0; j < v.size(); j++) {
            sum += m[j][v[j]];
         }
         vector<int> temp(v.begin(), v.end());
         temp.insert(temp.begin(), sum);
         pq.push(temp);
         return;
      }
      v[i]++;
      if (!ok && v[i] < z)
      update(v, i + 1, true);
      v[i]--;
      update(v, i + 1, ok);
   }
   int kthSmallest(vector<vector<int> >& m, int k){
      this->m = m;
      int ret = 0;
      int n = m.size();
      z = m[0].size();
      for (int i = 0; i < n; i++) {
         ret += m[i][0];
      }
      vector<int> temp(n);
      temp.insert(temp.begin(), ret);
      pq.push(temp);
      set<vector<int> > s;
      while (k--) {
         vector<int> temp = pq.top();
         pq.pop();
         s.insert(temp);
         ret = temp[0];
         temp.erase(temp.begin());
         update(temp, 0);
         while (!pq.empty() && s.count(pq.top())) {
            pq.pop();
         }
      }
      return ret;
   }
};
main(){
   Solution ob;
   vector<vector<int>> v = {{1,3,11},{2,4,6}};
   cout << (ob.kthSmallest(v, 5));
}

입력

{{1,3,11},{2,4,6}}

출력

7

마무리

이 알고리즘은 모든 조합을 미리 생성하는 대신, 우선순위 큐를 이용해 작은 합부터 순차적으로 탐색하기 때문에 k가 작을 때 특히 효율적입니다. 또한 집합(set)을 활용해 이미 처리한 조합의 중복을 제거함으로써 불필요한 연산을 줄일 수 있습니다. 행마다 정렬된 행렬에서 K번째 최소 합을 구하는 유사한 문제(예: 리트코드 1439번 '정렬된 행렬의 K번째 최소 배열 합')에도 동일한 접근 방식을 적용할 수 있습니다.