求子序列的第 k 小和

给定一个正整数数组 nums 和 一个非负数 k ,求出第 k 小的子序列和。

相同的不同子序列之和视为不同的和。同时空子序列的和记为 0 。例如对于数组 [1, 2, 3] ,它的子序列和的从小到大的排列是

[
    0 = [],
    1 = [1],
    2 = [2],
    3 = [3],
    3 = [1, 2],
    4 = [1, 3],
    5 = [2, 3],
    6 = [1, 2, 3]
]

Solution

在任意一个子序列中,每一个数组中的元素要么出现在这个子序列中,要么没有出现。一个数组的子序列构成的集合实际上就是它的幂集,假设这个数组有 len 个元素,那么我们可以用长度为 len 位的数唯一地确定这个幂集中的一个子序列:0 表示该位对应的元素没有出现在子序列中,1 则表示出现。所以,上面例子中的所有子序列可以表示为下面的数字:

[
    0 = [] = 0,
    1 = [1] = 1,
    2 = [2] = 2,
    3 = [3] = 4,
    3 = [1, 2] = 3,
    4 = [1, 3] = 5,
    5 = [2, 3] = 6,
    6 = [1, 2, 3] = 7
]

定义 S 为从表示子序列的数到子序列和的映射,问题转换为我们需要找到一个枚举子序列数 p 的方法满足:

满足第一个要求很简单:只需要将 S(p) 加入到一个小根堆中,使用堆排序。我们知道最小的子序列和一定是 0 对应的子序列 [] ,所以对于一个元素个数大于 0 的已从小到大排序的数组,除了空序列外,最小的子序列一定是 1 对应的 [n[0]]。所以我们只需要找到一个方法,从 1 出发,不重复地找到所有长度小于等于 len 位的数。

假设我们已经枚举了长度小于等于 l 的所有数,现在需要枚举所有长度等于 l + 1 的数。首先,对于长度为 l + 1 的数,其第 l 位一定是 1 ,第 l - 1 位要么是 0 要么是 1 ,这意味着对于一个长度等于 l + 1 的数 p 来说,总是存在唯一一个长度为 l 或 l - 1 的数 q 将第 l 位置一得到的,S(p) = S(q) + n[l] 。同时,这个长度为 l - 1 的数也可以通过一个长度为 l 的数通过将第 l - 1 位置零得到。所以,假设我们现在有一个长度为 l 的数 p ,我们可以求到两个长度为 l + 1 的数对应的 S :

S(q1) = S(p) + n[l]
S(q2) = S(p) + n[l] - n[l - 1]

因为每个数的来源是唯一的,所以通过这种方法枚举的数是不重复的,同时所有数都能枚举到(从个数来看),所以这个枚举方法是满足要求的。

fn perm_sub_sum(nums: &mut [i32]) -> Vec<i32> {
    use std::cmp::Reverse;
    use std::collections::BinaryHeap;

    let len = nums.len();
    let mut heap = BinaryHeap::new();
    let mut sums = vec![0];
    nums.sort_unstable();
    heap.push((Reverse(nums[0]), 0));

    while let Some((Reverse(sum), i)) = heap.pop() {
        sums.push(sum);
        if i < len - 1 {
            heap.push((Reverse(sum + nums[i + 1]), i + 1));
            heap.push((Reverse(sum + nums[i + 1] - nums[i]), i + 1));
        }
    }

    sums
}