算法基础算法

带约束条件的TopK选择问题，反悔贪心的思想

字数统计: 4.9k字   |   阅读时长: 20分
, ,

摘要: 堆贪心解决带约束条件的 TopK 问题

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

TopK 问题是算法中的一大类问题，很多问题可以抽象为 TopK 问题，然后用 TopK 问题的算法解决。在文章 topK问题：第K个最大元素 中，我们讨论了 TopK 问题的算法，在文章 topK问题分类汇总 中，我们汇总了 TopK 问题的题目清单。

TopK 问题是给定一组元素，找到其按从大到小排名第 K 的元素。其中一种算法是构造一个容量为 K 的最小堆，将所有元素压入一次，最终堆中保留的就是前 K 个最大元素，堆顶就是排名第 K 的元素。

有的时候元素的选择是有条件的，并且随着条件的放松，可选择的对象集合单调增加，此时还是用上述的方法，通过一个堆来维护当前可选择的对象，堆的容量限制为 K，这样堆中维护的始终是当前条件下，最优的 K 个元素。

具体做法是先将对象排序，越不满足约束条件越排在前面，然后枚举对象，压入堆中，若堆中元素个数超过 K，则弹出一个，直至对象集合耗尽。最终留在堆中的就是选出的 K 个对象。

还可以从反悔贪心的角度理解这种最优选择问题，参考文章 反悔贪心与反悔堆的思想，带截止时间的任务选择问题。

抽象问题：从 N 个对象中选出若干个

假设有 N 个对象，我们要从中选出若干个对象，满足一些约束条件，同时最优化某个指标。

枚举这 N 个对象，每轮枚举的过程如下：

  • step1: 放出一个对象（即当前枚举的对象）进入候选集；
  • step2: 根据约束条件和最优化指标，如果有必要，则从候选集淘汰一个（如果淘汰的是当前枚举的对象，则相当于跳过当前对象；如果淘汰的是此前枚举的对象，相当于反悔操作）。

候选集中的对象始终是取到局部最优解的选择，每次候选集出现变化，都相当于取到一个新的局部最优解，在这些取到的局部最优解中取一个最好的作为全局最优解的答案。但注意，全局最优解一定出现在这些局部最优解中这件事需要证明。

从反悔贪心的角度理解上述流程就是：一个对象如果满足进候选集的约束条件，就先让它进，后续当我们找到不满足进候选集的条件但能使得最优化指标变得更好的对象时，就放弃候选集中的某个对象使得当前对象可以进，放弃候选集中的对象就相当于一次反悔。

在放弃对象时，放弃那个对最优化的指标贡献最小的，这样就可以用堆来维护已选择的对象，对最优化指标贡献最小的放在堆顶，构成最小堆。

这就是用堆维护反悔贪心的场景，根据约束条件和最优化指标的不同，会有很多种不同的问题，下面我们来看一下 TopK 选择的问题。大致描述如下：

有 N 个对象，每个对象有两个指标，分别为指标1和指标2，从中选 K 个，最优化某个目标。而目标同时取决于指标1和指标2。

一般的思路是把某个指标作为约束条件，随着该指标逐渐放松，有可能进入候选集的对象逐渐增加。然后在堆中按照另一个指标维护已选择对象，根据最优化与另一指标的单调关系，决定用最大堆还是最小堆。如果需要反悔操作，就放弃堆顶的对象即可。

下面我们看两个类似的问题。

857. 雇佣 K 名工人的最低成本

有 n 名工人。 给定两个数组 quality 和 wage ，其中，quality[i] 表示第 i 名工人的工作质量，其最低期望工资为 wage[i] 。

现在我们想雇佣 k 名工人组成一个工资组。在雇佣 一组 k 名工人时，我们必须按照下述规则向他们支付工资：

  1. 对工资组中的每名工人，应当按其工作质量与同组其他工人的工作质量的比例来支付工资。
  2. 工资组中的每名工人至少应当得到他们的最低期望工资。

给定整数 k ，返回 组成满足上述条件的付费群体所需的最小金额 。在实际答案的 1e-5 以内的答案将被接受。。

提示：

1
2
3
n == quality.length == wage.length
1 <= k <= n <= 1e4
1 <= quality[i], wage[i] <= 1e4

示例 1：
输入： quality = [10,20,5], wage = [70,50,30], k = 2
输出： 105.00000
解释： 我们向 0 号工人支付 70，向 2 号工人支付 35。

示例 2：
输入： quality = [3,1,10,10,1], wage = [4,8,2,2,7], k = 3
输出： 30.66667
解释： 我们向 0 号工人支付 4，向 2 号和 3 号分别支付 13.33333。

算法：反悔贪心

K 个人分别是 $i_{1}, i_{2}, …, i_{K}$。

由于是按照工作质量同比例支付工资，因此总工资为 $(q[i_{1}] + q[i_{2}] + … + q[i_{k}]) x$。比例系数为 $x$。

由于需要满足最低期望工资，因此 $q[i_{k}]x \geq w[i_{k}]$，即 $x \geq w[i_{k}] / q[i_{k}]$。

因此取 $x = \min\limits_{1\leq j \leq K}\frac{w[i_{j}]}{q[i_{j}]}$ 即可，可以理解为性价比最高的人拿 $w[i]$ 的钱，其余的人根据工作质量按比例拿钱。

如果比例系数为 $x$，那么只有当 $w[i] / q[i] \leq x$ 时，人员 $i$ 才能进入工资组。也就是当 $x$ 越来越大时，可选的人员集合单调增加。

因此可以所有人按照 w[i] / q[i] 从小到大排序，然后依次枚举 $i$，对应的 $x = \frac{w[i]}{q[i]}$。此时从候选集中选出 $q$ 最小的 $K$ 个人，即为比例系数为 $x$ 时的 TopK 结果。

对象池中有 N 名工人，每个工人有一个最低期望工资 w 以及工作质量 q：

1
2
3
4
5
struct Worker
{
    int q, w;
    Worker(int q, int w):q(q),w(w){}
};

按照性价比 w/q （单位工作质量的期望工资）从小到大枚举工人，这样候选集中工人的单位工作质量的期望工资均小于等于当前枚举的工人。(这一步比较难，需要前面的推导)：

1
2
3
4
5
6
7
8
struct Cmp
{
    // 按 w[i] / q[i] 从大到小排序
    bool operator()(const Worker& worker1, const Worker& worker2) const
    {
        return (double)worker1.w / worker1.q < (double)worker2.w / worker2.q;
    }
};

首先将当前枚举的工人item压入堆中，相当于选了当前工人，当前已选工人的工作质量增加 item.q，如果堆中元素为 $K+1$，需要放弃一个工作质量最大的工人使人数变为 $K$ 个，这步放弃相当于反悔操作。

这样堆中的比较规则如下（最大堆）：

1
2
3
4
5
6
7
struct HeapCmp
{
    bool operator()(const Worker& worker1, const Worker& worker2) const
    {
        return worker1.q < worker2.q;
    }
};

代码 (C++)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
struct Worker
{
    int q, w;
    Worker(int q, int w):q(q),w(w){}
};

// k 个人分别是 i1, i2, ..., ik
// 总工资 = (q[i1] + q[i2] + ... + q[ik]) * x
// x 为常数且为 w[i] / q[i] 之一
//
// 若 i 选了，则 x[i] >= w[i] / q[i]
// 因此随着 w[i] / q[i] 越来越大，满足的候选越来越多
//
// 若选定了 x，即当前枚举到的 x(当前所有候选中最大)
// 选择所有候选中 q 最小的 k 个

struct Cmp
{
    // 按 w[i] / q[i] 从大到小排序
    bool operator()(const Worker& worker1, const Worker& worker2) const
    {
        return (double)worker1.w / worker1.q < (double)worker2.w / worker2.q;
    }
};

struct HeapCmp
{
    bool operator()(const Worker& worker1, const Worker& worker2) const
    {
        return worker1.q < worker2.q;
    }
};

class Solution {
public:
    // 若招的人为
    double mincostToHireWorkers(vector<int>& quality, vector<int>& wage, int K) {
        int n = quality.size();
        vector<Worker> workers;
        for(int i = 0; i < n; ++i)
            workers.emplace_back(quality[i], wage[i]);
        sort(workers.begin(), workers.end(), Cmp());
        priority_queue<Worker, vector<Worker>, HeapCmp> pq;
        int sum_q = 0;
        for(int i = 0; i < K; ++i)
        {
            sum_q += workers[i].q;
            pq.push(workers[i]);
        }
        double result = sum_q * (double)workers[K - 1].w / workers[K - 1].q;
        for(int i = K; i < n; ++i)
        {
            if(pq.top().q > workers[i].q)
            {
                sum_q -= pq.top().q;
                pq.pop();
                pq.push(workers[i]);
                sum_q += workers[i].q;
                result = min(result, sum_q * (double)workers[i].w / workers[i].q);
            }
        }
        return result;
    }
};

1383. 最大的团队表现值

公司有编号为 1 到 n 的 n 个工程师，给你两个数组 speed 和 efficiency ，其中 speed[i] 和 efficiency[i] 分别代表第 i 位工程师的速度和效率。请你返回由最多 k 个工程师组成的 ​​​​​​最大团队表现值 ，由于答案可能很大，请你返回结果对 10^9 + 7 取余后的结果。

团队表现值 的定义为：一个团队中「所有工程师速度的和」乘以他们「效率值中的最小值」。

提示：

1
2
3
4
5
6
1 <= n <= 10^5
speed.length == n
efficiency.length == n
1 <= speed[i] <= 10^5
1 <= efficiency[i] <= 10^8
1 <= k <= n

示例 1：
输入：n = 6, speed = [2,10,3,1,5,8], efficiency = [5,4,3,9,7,2], k = 2
输出：60
解释：
我们选择工程师 2（speed=10 且 efficiency=4）和工程师 5（speed=5 且 efficiency=7）。他们的团队表现值为 performance = (10 + 5) * min(4, 7) = 60 。

示例 2：
输入：n = 6, speed = [2,10,3,1,5,8], efficiency = [5,4,3,9,7,2], k = 3
输出：68
解释：
此示例与第一个示例相同，除了 k = 3 。我们可以选择工程师 1 ，工程师 2 和工程师 5 得到最大的团队表现值。表现值为 performance = (2 + 10 + 5) * min(5, 4, 7) = 68 。

示例 3：
输入：n = 6, speed = [2,10,3,1,5,8], efficiency = [5,4,3,9,7,2], k = 4
输出：72

算法：反悔贪心

最优化速度和和效率最小值的乘积。变化的量有两个，一个是速度，一个是效率。

其中 efficiency 取的是最小值，当 efficiency 取值确定的时候，只有工程师的效率值不小于 efficiency 时才能作为候选，随着 efficiency 越来越小，候选集合单调增加。

当固定 efficiency 时候，目标是从满足条件的 speed 中选择最大的 k 个，其中只要工程师的 efficiency 不小于当前枚举的 efficiency 时即满足条件。

对象池是 N 个工程师，每个工程师有一个速度 speed 和一个效率 efficiency：

1
2
3
4
5
struct Engineer
{
    int speed, efficiency;
    Engineer(int s, int e):speed(s),efficiency(e){}
};

按照效率从大到小枚举工程师，这样候选集中的工程师的效率均大于当前枚举的工程师：

1
2
3
4
5
6
7
8
9
struct Cmp
{
    bool operator()(const Engineer& engineer1, const Engineer& engineer2)
    {
        if(engineer1.efficiency == engineer2.efficiency)
            return engineer1.speed > engineer2.speed;
        return engineer1.efficiency > engineer2.efficiency;
    }
};

首先将当前枚举的工程师item压入堆中，相当于选了当前工程师，当前已选工程师的速度增加 item.speed，如果堆中元素为 $K+1$，需要放弃一个速度最小的工程师使人数变为 $K$，这步放弃相当于反悔操作。

由于是最多选 K 个人而不是恰好选 K 个人，因此堆中不用先初始化好 K 个元素。

代码 (C++)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
struct Engineer
{
    int speed, efficiency;
    Engineer(int s, int e):speed(s),efficiency(e){}
};

struct Cmp
{
    bool operator()(const Engineer& engineer1, const Engineer& engineer2)
    {
        if(engineer1.efficiency == engineer2.efficiency)
            return engineer1.speed > engineer2.speed;
        return engineer1.efficiency > engineer2.efficiency;
    }
};

class Solution {
public:
    int maxPerformance(int n, vector<int>& speed, vector<int>& efficiency, int k) {
        vector<Engineer> engineers;
        for(int i = 0; i < n; ++i)
            engineers.push_back(Engineer(speed[i], efficiency[i]));
        sort(engineers.begin(), engineers.end(), Cmp());
        ll result = 0;
        ll sum = 0;
        int min_efficiency = INT_MAX;
        priority_queue<int, vector<int>, greater<int> > pq;
        for(int i = 0; i < n; ++i)
        {
            min_efficiency = min(min_efficiency, engineers[i].efficiency);
            pq.push(engineers[i].speed);
            sum += engineers[i].speed;
            if((int)pq.size() > k)
            {
                sum -= pq.top();
                pq.pop();
            }
            result = max(result, sum * min_efficiency);
        }
        return result % MOD;
    }

private:
    using ll = long long;
    const int MOD = 1e9 + 7;
};

2813. 子序列最大优雅度

给你一个长度为 n 的二维整数数组 items 和一个整数 k 。

items[i] = [profiti, categoryi]，其中 profiti 和 categoryi 分别表示第 i 个项目的利润和类别。

现定义 items 的 子序列 的 优雅度 可以用 total_profit + distinct_categories^2 计算，其中 total_profit 是子序列中所有项目的利润总和，distinct_categories 是所选子序列所含的所有类别中不同类别的数量。

你的任务是从 items 所有长度为 k 的子序列中，找出 最大优雅度 。

用整数形式表示并返回 items 中所有长度恰好为 k 的子序列的最大优雅度。

注意：数组的子序列是经由原数组删除一些元素（可能不删除）而产生的新数组，且删除不改变其余元素相对顺序。

提示：

1
2
3
4
5
6
7
1 <= items.length == n <= 1e5
items[i].length == 2
items[i][0] == profiti
items[i][1] == categoryi
1 <= profiti <= 1e9
1 <= categoryi <= n 
1 <= k <= n

示例 1：
输入：items = [[3,2],[5,1],[10,1]], k = 2
输出：17
解释：
在这个例子中，我们需要选出长度为 2 的子序列。
其中一种方案是 items[0] = [3,2] 和 items[2] = [10,1] 。
子序列的总利润为 3 + 10 = 13 ，子序列包含 2 种不同类别 [2,1] 。
因此，优雅度为 13 + 22 = 17 ，可以证明 17 是可以获得的最大优雅度。

示例 2：
输入：items = [[3,1],[3,1],[2,2],[5,3]], k = 3
输出：19
解释：
在这个例子中，我们需要选出长度为 3 的子序列。
其中一种方案是 items[0] = [3,1] ，items[2] = [2,2] 和 items[3] = [5,3] 。
子序列的总利润为 3 + 2 + 5 = 10 ，子序列包含 3 种不同类别 [1, 2, 3] 。
因此，优雅度为 10 + 32 = 19 ，可以证明 19 是可以获得的最大优雅度。

示例 3：
输入：items = [[1,1],[2,1],[3,1]], k = 3
输出：7
解释：
在这个例子中，我们需要选出长度为 3 的子序列。
我们需要选中所有项目。
子序列的总利润为 1 + 2 + 3 = 6，子序列包含 1 种不同类别 [1] 。
因此，最大优雅度为 6 + 12 = 7 。

算法：反悔贪心

设选出的 $k$ 个对象的利润分为为 $p_{1}, \cdots, p_{k}$，类别为 $c_{1}, \cdots, c_{k}$。则最优化目标为：

其中 $C(c_{1}, \cdots, c_{k})$ 表示 $k$ 个对象的种类数。

先将利润最大的 K 歌对象压入候选集，得到一个目标值，该目标值可能是全局最优的，也可能不是。因为此时放一个利润更低的新对象进来，替换掉候选集中的某个已选对象，虽然候选集总利润变低，但种类数可能增加，替换后可能使得结果变得更好。这步替换相当于反悔操作。

对象池中有 N 个对象，每个对象有一个利润 p 和一个种类 c：

1
2
3
4
5
6
7
struct Item
{
    int p;
    int c;
    Item(){}
    Item(int p, int c):p(p),c(c){}
};

按照利润从大到小枚举对象，这样候选集中的对象的利润均大于当前枚举的对象：

1
2
3
4
5
6
7
struct Cmp
{
    bool operator()(const Item& i1, const Item& i2) const
    {
        return i1.p > i2.p;
    }
};

如果候选集中元素个数小于 $K$，则直接压入即可。如果候选集中元素个数等于 $K$，则需要判断是跳过当前对象，还是放弃候选集中某个对象，然后将当前对象压入。

设当前枚举的对象为 item，此时候选集中对象的利润都大于等于 item.p，候选集中的总利润为 P，候选集中的种类数为 C，P + C * C 即为在种类数为 C 的情况下的局部最优解。下面我们要找的是当 C 变得更大时的局部最优解。

若 item.c 与候选集中某个对象的种类相同，则不论用 item 替换谁都不会使得种类数 C 变大，直接跳过即可。

若 item.c 与候选集中任意一个对象都不一样，则在候选集找到一个出现过多次的对象，这样将找到的对象替换为 item 后，种类数会增加 1。

如果候选集中不存在出现多次的对象，则直接跳过 item；如果候选集中存在出现多次的对象，则选出其中利润最小的那一个，记其为 t，执行替换，将 t 替换为 item，此时候选集中为种类数为 C + 1 情况下的局部最优解，用它更新全局最优解的答案。

综上，在反悔操作时，我们要放弃的是候选集中类别有重复且利润最小的对象。可以用最小堆维护候选集中的对象，使得堆顶就是利润最小的那一个。这里还需要注意堆中的对象必须是在候选集中重复的，可以用一个哈希集合来保存候选集中出现过的类别。这样在按利润从大到小枚举对象时，若其类别没出现过，则插入哈希表而不插入堆；若出现过（哈希表中有记录），则插入到堆中。

对于本题，注意到对象的枚举是按照利润从大到小来的，因此也可以用栈代替堆，这样栈顶自然就是利润最小的。压入、弹出的时间复杂度比使用堆更好。但是将对象排序依然要 $O(n\log n)$，所以总体时间复杂度还是 $O(n\log n)$。

代码 (C++)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
struct Item
{
    int p;
    int c;
    Item(){}
    Item(int p, int c):p(p),c(c){}
};

struct Cmp
{
    bool operator()(const Item& i1, const Item& i2) const
    {
        return i1.p > i2.p;
    }
};

class Solution {
public:
    long long findMaximumElegance(vector<vector<int>>& items, int k) {
        int n = items.size();
        vector<Item> _items(n);
        for(int i = 0; i < n; ++i)
        {
            _items[i].p = items[i][0];
            _items[i].c = items[i][1];
        }
        sort(_items.begin(), _items.end(), Cmp());
        unordered_set<int> setting;
        stack<int> st;
        long long P = 0; // 候选集中的总利润
        for(int i = 0; i < k; ++i)
        {
            P += _items[i].p;
            if(setting.count(_items[i].c) == 0)
            {
                setting.insert(_items[i].c);
            }
            else
            {
                st.push(_items[i].p);
            }
        }
        long long C = setting.size(); // 候选集中的种类数
        long long ans = P + C * C;
        for(int i = k; i < n; ++i)
        {
            if(setting.count(_items[i].c) > 0 or st.empty())
                continue;
            P += _items[i].p - st.top();
            st.pop();
            setting.insert(_items[i].c);
            C = setting.size();
            ans = max(ans, P + C * C);
        }
        return ans;
    }
};
Share