0629：K 个逆序对数组（★★）

力扣第 629 题

题目

对于一个整数数组 nums，逆序对是一对满足 0 <= i < j < nums.length 且 nums[i] > nums[j]的整数对 [i, j] 。

给你两个整数 n 和 k，找出所有包含从 1 到 n 的数字，且恰好拥有 k 个 逆序对 的不同的数组的个数。由于答案可能很大，只需要返回对 109 + 7 取余的结果。

示例 1：

输入：n = 3, k = 0
输出：1
解释：
只有数组 [1,2,3] 包含了从1到3的整数并且正好拥有 0 个逆序对。

示例 2：

输入：n = 3, k = 1
输出：2
解释：
数组 [1,3,2] 和 [2,1,3] 都有 1 个逆序对。

提示：

• 1 <= n <= 1000
• 0 <= k <= 1000

分析

假设 n 放到位置 idx，那么包含 n 的逆序对有 n-idx-1，然后转为了子问题：1 到 n-1 的排列中恰好 k-(n-idx-1) 个逆序对的个数。

因此令 dp[i][j] 代表 1 到 i 的排列中恰好 j 个逆序对的个数，即可递推：

dp[i][j] = sum(dp[i-1][j-cnt] for cnt in range(i) if j-cnt>=0)

但这个时间复杂度会超时。

观察发现递推式中其实是 dp[i-1] 的连续区间的和，于是想到用前缀和。 事先保存 dp[i-1] 的前缀和数组 pre，递推式变为：

dp[i][j] = pre[j+1]-pre[max(0, j-i+1)]

还可以用滚动数组将 dp 优化为一维数组。

解答

1
2
3
4
5
6

def kInversePairs(self, n: int, k: int) -> int:
    dp, mod = [1]+[0]*k, 10**9+7
    for i in range(1, n+1):
        pre = list(accumulate([0]+dp, lambda x,y:(x+y)%mod))
        dp = [pre[j+1]-pre[max(0, j-i+1)] for j in range(k+1)]
    return dp[-1]%mod

548 ms