「技巧总结」第k小和贪心问题

✈

基础问题描述

给定 $n$ 个数，求选出一个子集，和的第 k 大是多少

$n \le 10^6, k \le 10^6$

思维过程

首先我们可以想到一个很暴力的贪心，最开始的时候我们往堆里加一个空集，然后每次从堆里拿出当前最小的子集，考虑用这个子集状态转移到一个新的状态。

显然如果这种转移满足一下条件：

1. 可以不重不漏的考虑到每一种状态。

   本质上是每种状态都能被拿出恰好 1 次

2. 第 k 个被拿出的状态一定是第 k 大的（一个状态被拿出时保证每个比他小的状态都被拿出了）

   这个限制本质就是小的状态只能转移到比他大的状态，因为所有状态都会被转移到，所以就能满足第 k 大是第 k 个被拿出的

那么显然我们只要拿出 k 次堆顶，那么最后一次的堆顶就是第 k 大的值。

假设当前状态 $mask$ 最后一个选的数为 $i$ (如果是空集的话 $i = 0$) ，那么暴力的转移就是枚举一个 $j>i$ ，把 mask+(1<<j) 加入堆，那么显然满足上面的条件。

但是时间复杂度显然不对，因为一个数会有 $n$ 种转移，所以这个做法是 $O(nm)$ 的。

考虑如何优化这个转移。

我们想到了令一个状态，原本我们的状态是一个 $mask$ 表示当前的子集，现在我们把状态改成 $(mask,i)$ ，表示当前只考虑了前 $i$ 个数，状态为 $mask$ ，转移有两种，和一般的背包一样：每次拿当前的物品或者不拿，然后让 i++ ，也就是跳到下一个物品，这种状态转移减少了，但是注意到很多状态的 mask 是一样的，而每种 mask 只能贡献一次答案，所以这样做的话有很多中间状态是无用的，也就意味着我们拿出的第 k 个数并不是第 k 大，而可能是第 $k-x$ 大$(0 \le x < k)$ 。

考虑把这两个状态合并一下，我们的状态是 $(mask,i)$ 表示当前只考虑了前 $i - 1$ 个数，强制第 i 个数必选，转移的时候按照第二种状态转移，就是考虑选或不选第 i 个数，然后 ++i ，由强制选第 i 个数变成强制选第 i + 1 个数。如果特判掉空集，这个状态可以转移出所有状态，显然不满足第二种限制，有可能会出现强制选了 i + 1 后而不选择 i 使总和变小了。但是我们发现如果我们保证第 i + 1 个物品总是大于第 i 个物品（也就是把物品排序），那么转移就一定是从小往大的了。至此，我们的复杂度降到了 $O(m\log m)$ ，至于空集简单特判即可。

extra I

给定 n 个集合，每个集合包含 $s_i$ 个数，第 j 个数为 $a_{i,j}$，现在要从每个集合里拿出了一个数，他的代价是每个数的和，求第 k 大的代价是多少。

$\sum s_i \le 10^6,n\le10^6,1 \le s_i$

Solution

首先我们把每个按照我们的思路，首先将每个集合里的数排序，然后设计状态 $(mask,i)$ 表示现在考虑了前 i - 1 个集合，并且强制第 $i$ 个集合选的是第 2 个数，转移也比较ez，每次考虑 $i - 1$ 这个集合是否选一个更大的数，或者跳过这个集合转向下一个集合。特判掉大小为 1 的集合后，我们把每个集合按照 $a_{i,2}-a_{i,1}$ 排序，那么每次强制 i 集合选第二个数变为强制选 i + 1 选第二个数的时候，代价一定是增加的。唯一的问题就是因为我们强制了 i + 1 个集合只能选第二个数，无法达到每个状态。这个很容易解决，我们设计一种特殊状态：$(mask,i,j)$ 表示当前考虑了前 i - 1 个集合，强制第 i 集合选第 j + 1 个数，我们强制只有 $(mask,i,1)$ 能转移到 $(mask,i+1,1)$，剩下的状态只能转移 i - 1 这个集合选更大的数或是 $(mask,i,j+1)$ ，那么转移种数还是 $O(1)$ 并且满足上面的两个条件。还是特判一下每个集合都只选第一个数的情况。

时间复杂度：$O(m\log m)$

*更普遍的做法

之前做过一个类似的题，对于这类问题，我们有一种更通用、更简单的做法，不过时间复杂度稍劣，一般是 $O(n\log n \log V)$

以基础问题为例，一般看到这类第 k 大问题，我们除了会想到上面的那个贪心，还会想到另一个算法：二分答案。但是二分完答案后如何 check 一个答案是否合法呢？有一个更暴力的想法，我们直接把物品按照任意顺序 排列，然后从前往后做背包，用一个堆来记录用前 i 个物品能组成哪些权值小于我们二分的 mid 的背包。每次加一个物品的时候，从堆里从小往大拿出每个一方案，看看能否将第 $i$ 个物品放入，如果做到某一步的时候发现当前的背包价值已经无法加入第 $i$ 个物品就停止（显然后面的方案都放不下了），一直坐过去，如果在加入某一个物品的时候发现当前合法的总方案数已经超过 $k$ 种，那么我们二分的这个 mid 一定大于等于答案，否则说明当前的 mid 小于答案。考虑这么做的复杂度是什么，每次 check 我们只要发现方案数大于 k 就会直接退出，所以最多加入 k 物品，最多从堆里查询 k + n 个物品，所以只要精细的实现一下， check 复杂度是 $O((k + n) \log k)$ 的，总复杂度如上所述 $O(n \log n \log V)$。

加入物品的代码大致如下

vector <int> tmp; // 存储新加的物品
bool add(min_heap p, int q)
{
    tmp.clear();
    for(int i : p)
    {
        if(i + q > mid) break;
        tmp.pb(i + q);
        if(tmp.size() + p.size() >= k) return 1;
    }
    for(int i : tmp) p.push(i);
    return 0;
}

extra I 解法二

我们按照上面的思路，还是考虑每次加入一个物品，但是由于每个集合共有 $s_i - 1$ 种选择，所以我们只要上面的代码稍微改一改就好了。

vector <int> tmp; // 存储新加的物品
bool add(min_heap p, int q[], int s_i)
{
    tmp.clear();
    for(int j = 1; j < s_i; ++ j)
    {
  		for(int i : p)
    	{
       		if(i + q[j] > mid) break;
       		tmp.pb(i + q[j]);
        	if(tmp.size() + p.size() >= k) return 1;
    	}
    }
    for(int i : tmp) p.push(i);
    return 0;
}

extra II

给定一棵树，点有点权，求树上第 k 大联通块。

$n \le 10^5,k\le 10^5$

二分之后直接 check，不妨设 1 为根，把一个联通块在深度最浅的那个点考虑。令 $p_i$ 表示一个堆，维护所有最浅点为 i 的合法联通块。考虑如何求出 $p_u$ ，只要我们求出了他的子树的 $p_v$，然后把两个堆合并起来就可以了，和上面大同小异，也是做到 k 个就不做了。

只要把 extra I 合并代码的数组改成另一个堆就即可。