UVa OJ 11997 - K Smallest Sums

https://uva.onlinejudge.org/index.php?option=com_onlinejudge&Itemid=8&page=show_problem&problem=3148

K <= 750

K개의 배열이 있고 각 배열은 K개의 원소가 있다. 

각 배열에서 한개의 수를 골라 K개의 수를 골라 모두 더할 것인데, 이 합이 K번째로 작은 합을 구하는 문제이다.

우선 naive 하게 풀어본다면 K개의 배열에서 각각 K개의 경우의 수가 있기 때문에 K^K 개의 경우의 수가 있으며

이것을 정렬하면 O(K^(K+1)lgK) 의 시간복잡도를 갖는다. 이 것을 어떻게 줄일 수가 있을까?

답부터 말하자면 우선 각 행을 각각 정렬을 한다. 그리고 위에서부터 한 행씩 보는데 i 번째 행까지 각 행의 수를 하나씩 선택 했을 때

가장 작은 합부터 K개를 계속 유지를 한다. 그리고 i+1번째 행을 보고 다시 지금까지 구한 가장 작은 K개의 합들과 잘 합쳐서

다시 i+1번째 행까지 봤을 때의 가장 작은 K개의 합을 유지하며 이렇게 K 번째 행까지 반복하면 답이 구해질 것이다.

그럼 과연 i 번째 행까지 각 행에서 한개의 수를 골랐을 때의 가장 작은 K개의 합과 i+1번째 행을 어떻게 합칠까?

일단 간단하게 생각하였을 때 현재 유지하고있는 K개의 합과 새로 보는 행의 K개의 수가 있기 때문에 K^2개의 새로운 합이 생기며

이것을 정렬하여 K개를 선택하는 방식을 취한다면 O(K^2lgK) 의 시간복잡도를 가질것이고 K개의 행을 볼 것이기 때문에

O(K^3lgK) 인데 K^3 은 대략 4억정도가 되는데 1초안에 실행되어야하기 때문에 이건 안된다. 그렇다면 이렇게 한번 생각해보자

지금까지구한 K개와 새로 더해질 K개에 대해 투포인터를 사용하는데 우선순위큐를 사용하여 현재의 두 포인터가 가리키는 수에서

각각 다음 숫자로 바꾸었을 때의 합과 두포인터의 묶음 2개를 우선순위큐에 넣어서 다음은 가장 작은 합을 가지는 포인터 쌍을 큐에서 빼서

이 합을 가장작은 K개의 합에 추가하고 같은 동작을 반복한다. 이렇게 하면 O(KlgK) 가 걸릴 것이고 K개의 행을보면 O(K^2lgK)의

시간복잡도가 될 것이다. 이정도면 충분한 시간 복잡도이지만 여기에는 큰 문제가 하나 있다. 

O O O O O O O

   ^

O O O O O O O

   ^

이렇게 두 포인터가 두번째를 가리키게 되는 상황으로 오는 방법이 두가지가 있다.  위에 있는 포인터가 한칸 전에 있을 때와 

아래에 있는 포인터가 한칸 전에 있을 때 이렇게 두가지에서 저 상황이 될 수 있기 때문에 중복된 원소가 각각 두개씩 계산이 된다.

그렇기 때문에 중복 문제만 해결하면 되는데 어떻게 해결할까.

한가지 생각해 볼 수 있는 방법은 위 포인터를 고정하여 아래 포인터만을 움직이는 식으로 하면 위 포인터는 총 K가지가 있기 때문에

K개의 포인터가 있다고 생각을 하는 것이다. 그렇다면 K개의 포인터가 어디서 시작하냐면

i 번째 포인터는 위쪽 K개 수에 대해선 i번 째 수로 고정하고 아래쪽 K개 수만 한칸 씩 다음 수로 바꾸어 선택하여 합을 구하면서

우선순위큐에 넣어주면 매 번 K개의 포인터 중 합이 가장 작은 합을 O(lgK) 로 선택할 수 있을 것이다. 이런식으로

매 번 K개의 포인터중 합이 가장작은 것을 새로뽑는 K개의 수에 추가하고 아래 포인터를 한칸옮겨 구한 합을 큐에 다시 새로넣어주고

이것을 반복하면 O(KlgK) 에 중복없이 두 K개의 배열을 합쳐 다시 가장 작은 K개의 합을 계산할 수 있다.

#include <bits/stdc++.h>
using namespace std;
const int MOD = 1e9+7;
const int INF = 0x3c3c3c3c;
const long long INFL = 0x3c3c3c3c3c3c3c3c;
 
int T,K;
int A[753][753];
long long knum[753];
int main() {
    while(scanf("%d",&K) != -1){
        memset(knum, 0, sizeof(0));
        for(int i = 0 ; i < K; i++) for(int j = 0 ; j < K; j++) scanf("%d",&A[i][j]);
        for(int i = 0 ; i < K; i++) sort(A[i], A[i]+K);
        for(int i = 0 ; i < K; i++) knum[i] = A[0][i];
        for(int i = 1 ; i < K; i++){
            priority_queue<pair<long long, int> > pq;
            for(int j = 0 ; j < K; j++) pq.push({-(knum[j] + A[i][0]), 0});
            for(int j = 0 ; j < K; j++) {
                auto p = pq.top();
                pq.pop();
                long long sum = -p.first;
                int idx = p.second;
                knum[j] = sum;
                if(idx < K-1) pq.push({-(sum - A[i][idx] + A[i][idx+1]), idx+1});
            }
        }
        for(int i = 0 ; i < K; i++) printf("%lld",knum[i]), printf("%c"," \n"[i==K-1]);
    }
    return 0;
}