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

C++로 풀어보는 K 역쌍(K Inverse Pairs) 배열 문제

두 정수 n과 k가 주어졌을 때, 1부터 n까지의 숫자로 구성된 배열 중에서 정확히 k개의 역쌍(inverse pair)을 가지는 서로 다른 배열의 개수를 구하는 문제입니다.

여기서 역쌍이란 배열 내 i번째 원소와 j번째 원소에 대해 i < j이면서 a[i] > a[j]인 경우를 의미합니다. 즉, 앞에 있는 숫자가 뒤에 있는 숫자보다 큰 경우가 역쌍이 됩니다.

정답은 매우 커질 수 있으므로, 결과는 $10^{9}$ + 7로 나눈 나머지를 출력해야 합니다.

예제 이해하기

예를 들어 입력이 n = 3, k = 1이라면 출력은 2가 됩니다. 그 이유는 [1, 3, 2]와 [2, 1, 3] 두 개의 배열만이 정확히 하나의 역쌍을 가지기 때문입니다.

접근 방법: 동적 계획법(DP)

이 문제는 동적 계획법을 활용하여 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • (n + 1) × (k + 1) 크기의 2차원 배열 dp를 선언합니다.
  • dp[i][j]는 1부터 i까지의 숫자로 만든 배열 중 정확히 j개의 역쌍을 가지는 배열의 개수를 나타냅니다.

알고리즘 단계

  1. dp[0][0] := 1로 초기화합니다.
  2. i를 1부터 n까지 반복하면서:
    • dp[i][0] := 1로 설정합니다 (역쌍이 없는 경우는 오름차순 배열 하나뿐).
    • j를 1부터 k까지 반복하면서:
      • dp[i][j] := dp[i][j - 1] + dp[i - 1][j]
      • dp[i][j] := dp[i][j] mod m
      • 만약 j ≥ i라면:
        • dp[i][j] := (dp[i][j] - dp[i - 1][j - i] + m) mod m
  3. 최종적으로 dp[n][k]를 반환합니다.

점화식에서 dp[i][j - 1]을 더하고 dp[i - 1][j - i]를 빼는 이유는, 새 숫자 i를 삽입할 때 만들 수 있는 역쌍의 범위를 제한하여 중복 계산을 방지하는 접두사 합(prefix sum) 최적화 기법 때문입니다. m을 더해주는 부분은 음수가 되는 것을 방지하기 위한 처리입니다.

C++ 구현 예제

#include <bits/stdc++.h>
using namespace std;
const int m = 1e9 + 7;
class Solution {
public:
    int kInversePairs(int n, int k) {
        vector<vector<int>> dp(n + 1, vector<int>(k + 1));
        dp[0][0] = 1;
        for(int i = 1; i <= n; i++){
            dp[i][0] = 1;
            for(int j = 1; j <= k; j++){
                dp[i][j] = dp[i][j - 1] + dp[i - 1][j];
                dp[i][j] %= m;
                if(j >= i){
                    dp[i][j] = (dp[i][j] - dp[i - 1][j - i] + m) % m;
                }
            }
        }
        return dp[n][k];
    }
};
main(){
    Solution ob;
    cout << (ob.kInversePairs(4, 2));
}

입력

4
2

출력

5

결과 분석

n = 4, k = 2일 때 출력값은 5입니다. 실제로 [1, 4, 2, 3], [1, 3, 4, 2], [2, 1, 4, 3], [2, 3, 1, 4], [3, 1, 2, 4] 등 총 5개의 배열이 정확히 2개의 역쌍을 가지므로 올바른 결과임을 확인할 수 있습니다.

이 알고리즘의 시간 복잡도는 O(n × k), 공간 복잡도 역시 O(n × k)입니다. 접두사 합 기법을 사용하지 않으면 시간 복잡도가 O(n × k × min(n, k))까지 늘어날 수 있으므로, 위와 같은 최적화가 성능 면에서 매우 중요합니다.