LeetCodeの629. K Inverse Pairs Arrayという問題ですが、解答にあたりかなり苦労したので、解き方や考え方のメモを残しておきたいと思います。
この問題は「1~nからなる数列を考える。invert pairの数が k 個の数列を数えなさい」というものです。invert pairとは、$0 \leqq i < j < n$ かつ $nums[i] > nums[j]$ を満たす組のこと。たとえば[4, 2, 3, 1]という数列があるとすると、(4,2), (4,3), (4,1), (2,1), (3,1)がinvert pairになります。
まず具体例から考えます。1~4からなる数列として[4, 2, 3, 1]があります。この数列に5を挿入すると、invert pairの数がどのように増えるのかを観察してみましょう。
| もとの配列 | 挿入後の配列 | 増えたinvert pair | 増えたinvert pairの数 |
|---|---|---|---|
| [4, 2, 3, 1] | [5, 4, 2, 3, 1] | (5, 4), (5, 2), (5, 3), (5, 1) | 4 |
| [4, 2, 3, 1] | [4, 5, 2, 3, 1] | (5, 2), (5, 3), (5, 1) | 3 |
| [4, 2, 3, 1] | [4, 2, 5, 3, 1] | (5, 3), (5, 1) | 2 |
| [4, 2, 3, 1] | [4, 2, 3, 5, 1] | (5, 1) | 1 |
| [4, 2, 3, 1] | [4, 2, 3, 1, 5] | - | 0 |
つまり、1~4からなる数列に5を挿入すると、以下のいずれかに遷移するとわかります。
- invert pairが0個増えた、1~5からなる数列
- invert pairが1個増えた、1~5からなる数列
- invert pairが2個増えた、1~5からなる数列
- invert pairが3個増えた、1~5からなる数列
- invert pairが4個増えた、1~5からなる数列
これを一般してみましょう。dp[p][q]を「反転数がq個の1~pからなる数列の数」とします。この数列にp+1を挿入し、1~p+1からなる数列にすることを考えます。すると、dp[p][q]から次のp個の状態に遷移します。
- dp[p+1][q+0]
- dp[p+1][q+1]
- dp[p+1][q+2]
- ...
- dp[p+1][q+p-1]
- dp[p+1][q+p]
あとはこれをプログラムに落とし込めばよいです。問題の指示にあるとおり、1000000007でわったあまりを求めることを忘れないようにしましょう。計算量は$O(n^3)$ですが、$n \leqq 1000$とnが比較的小さいので、高速な言語を利用すれば、問題なく制限時間に間に合います。
main.rs
impl Solution {
pub fn k_inverse_pairs(n: i32, k: i32) -> i32 {
let n = n as usize;
let k = k as usize;
let mut dp = vec![vec![0; k+1]; n+1];
dp[1][0] = 1;
for p in 1..n {
for q in 0..k+1 {
for r in 0..p+1 {
if q+r <= k {
dp[p+1][q+r] += dp[p][q];
dp[p+1][q+r] %= 1000000000 + 7;
}
}
}
}
return dp[n][k];
}
}