【对顶队列】【中位数贪心】【前缀和】3086. 拾起 K 个 1 需要的最少行动次数-阿里云开发者社区

本文涉及知识点

对顶队列（栈）分类讨论

LeetCode3086. 拾起 K 个 1 需要的最少行动次数

给你一个下标从 0 开始的二进制数组 nums，其长度为 n ；另给你一个正整数 k 以及一个非负整数 maxChanges 。

灵茶山艾府在玩一个游戏，游戏的目标是让灵茶山艾府使用最少数量的行动次数从 nums 中拾起 k 个 1 。游戏开始时，灵茶山艾府可以选择数组 [0, n - 1] 范围内的任何索引index 站立。如果 nums[index] == 1 ，灵茶山艾府就会拾起一个 1 ，并且 nums[index] 变成0（这不算作一次行动）。之后，灵茶山艾府可以执行任意数量的行动（包括零次），在每次行动中灵茶山艾府必须恰好执行以下动作之一：

选择任意一个下标 j != index 且满足 nums[j] == 0 ，然后将 nums[j] 设置为 1 。这个动作最多可以执行 maxChanges 次。

选择任意两个相邻的下标 x 和 y（|x - y| == 1）且满足 nums[x] == 1, nums[y] == 0 ，然后交换它们的值（将 nums[y] = 1 和 nums[x] = 0）。如果 y == index，在这次行动后灵茶山艾府拾起一个 1 ，并且 nums[y] 变成 0 。

返回灵茶山艾府拾起恰好 k 个 1 所需的最少行动次数。

示例 1：

输入：nums = [1,1,0,0,0,1,1,0,0,1], k = 3, maxChanges = 1

输出：3

解释：如果游戏开始时灵茶山艾府在 index == 1 的位置上，按照以下步骤执行每个动作，他可以利用 3 次行动拾取 3 个 1 ：

游戏开始时灵茶山艾府拾取了一个 1 ，nums[1] 变成了 0。此时 nums 变为 [1,0,0,0,0,1,1,0,0,1] 。

选择 j == 2 并执行第一种类型的动作。nums 变为 [1,0,1,0,0,1,1,0,0,1]

选择 x == 2 和 y == 1 ，并执行第二种类型的动作。nums 变为 [1,1,0,0,0,1,1,0,0,1] 。由于 y == index，灵茶山艾府拾取了一个 1 ，nums 变为 [1,0,0,0,0,1,1,0,0,1] 。

选择 x == 0 和 y == 1 ，并执行第二种类型的动作。nums 变为 [0,1,0,0,0,1,1,0,0,1] 。由于 y == index，灵茶山艾府拾取了一个 1 ，nums 变为 [0,0,0,0,0,1,1,0,0,1] 。

请注意，灵茶山艾府也可能执行其他的 3 次行动序列达成拾取 3 个 1 。

示例 2：

输入：nums = [0,0,0,0], k = 2, maxChanges = 3

输出：4

解释：如果游戏开始时灵茶山艾府在 index == 0 的位置上，按照以下步骤执行每个动作，他可以利用 4 次行动拾取 2 个 1 ：

选择 j == 1 并执行第一种类型的动作。nums 变为 [0,1,0,0] 。

选择 x == 1 和 y == 0 ，并执行第二种类型的动作。nums 变为 [1,0,0,0] 。由于 y == index，灵茶山艾府拾起了一个 1 ，nums 变为 [0,0,0,0] 。

再次选择 j == 1 并执行第一种类型的动作。nums 变为 [0,1,0,0] 。

再次选择 x == 1 和 y == 0 ，并执行第二种类型的动作。nums 变为 [1,0,0,0] 。由于y == index，灵茶山艾府拾起了一个 1 ，nums 变为 [0,0,0,0] 。

提示：

2 <= n <= 10⁵

0 <= nums[i] <= 1

1 <= k <= 10⁵

0 <= maxChanges <= 10⁵

maxChanges + sum(nums) >= k

分类讨论

1，消耗0行动次数。nums[index]为1。

2，消耗1行动次数，nums[index-1]或nums[index+1]为1。

3，消耗2行动次数，一转换次数。nums[index-1]或nums[index+1]设置为1，再移动。

4，消耗2或更多行动次数。将其它的1移动过来。 ⟺ \iff⟺ 计算距离index 最近的m个一。

枚举index，时间复杂度O(n)。前三种分类，时间复杂度都是O(1)，如何将第四类操作的时间复杂度将为O(1)。

如果有转换次数，一定用分类三，不用分类四。故需要讨论分类四时，maxChanges 次数已经用完，故m = k - maxChange。用对顶队列分别记录下标小于index 的1下标，下标大于等于index的1下标。

程序流程

一，用队列升序收集1的下标。

二，利用对顶队列计算距离index最近的m个1。

三，分别处理情况一到四。注意：拾取次数不能超过k。m个1要扣除分类一和分类二的1。

对顶队列预处理最近m个1

一，计算vPre[0]。

二，如果队列不发生变化，vPre[i2] = vPre[i2 - 1] + que1.size() - que2.size();

三，如果que2移到que1，触发条件que2.front() < i2。由于i2一次加1，所以que2.front()此时一定等于i2-1。应该是1，误算成-1。故+2。 que1入队，que2出队。

三，枚举i2是，同时枚举que1Index。index1.front()一定> i2，否则ndex1.front() < i2-1，应该进入que1。index是升序，故que1中的下标都小于ndex1.front()，故ndex1.front()比que1中的数距离i2-1都近。一定会淘汰que1中的数据。

因为ndex1.front()大于que2中的下标，故不会淘汰que2中的数。

que1中的数，越小越容易淘汰，故从队首开始淘汰。

代码

核心代码

class Solution {
public:
  long long minimumMoves(const vector<int>& nums, const int k, const int maxChanges) {
    m_c = nums.size();
    queue<int> que1Index;
    for (int i = 0; i < m_c; i++)
    {
      if (nums[i])
      {
        que1Index.emplace(i);
      }
    }
    vector<long long> vPre(nums.size());//距离 nums[i]最近的iNum个1的距离
    const int i1Num = k - maxChanges;   
    queue<int> que1, que2;
    while(que1Index.size() && (que2.size() < i1Num))
    {//计算vPre[0]
      que2.emplace(que1Index.front());
      vPre[0] += que1Index.front();
      que1Index.pop();
    }
    for (int i2 = 1; i2 < m_c; i2++)
    {
      vPre[i2] = vPre[i2 - 1] + que1.size() - que2.size();
      if (que2.size() && (que2.front() < i2))
      {//队列二向队列一移动
        que1.emplace(que2.front());
        que2.pop();
        vPre[i2] += 1 + 1;
      }
      while (que1.size() && que1Index.size() && (i2 - que1.front() > que1Index.front() - i2))
      {
        vPre[i2] -= (i2 - que1.front());
        vPre[i2] += que1Index.front() - i2;
        que2.emplace(que1Index.front());
        que1.pop();
        que1Index.pop();
      }
    }
    long long llRet = LLONG_MAX;;
    for (int i = 0; i < nums.size(); i++)
    {
      int iNeed = k - nums[i];  
      long long llDo = 0;
      if ((iNeed > 0) && (i > 0) && nums[i - 1])
      {
        iNeed--; llDo++;
      }
      if ((iNeed > 0) && (i+1 < nums.size()) && nums[i + 1])
      {
        iNeed--; llDo++;
      }
      const int iNeiBo = llDo;
      const int iChange = min(maxChanges, iNeed);
      iNeed -= iChange;
      llDo += 2 * iChange;
      if (iNeed > 0)
      {
        llDo += vPre[i] - iNeiBo;
      }
      llRet = min(llRet, llDo);
    }
    return llRet;
  }
  int m_c;
};

测试用例

template<class T, class T2>
void Assert(const T& t1, const T2& t2)
{
  assert(t1 == t2);
}
template<class T>
void Assert(const vector<T>& v1, const vector<T>& v2)
{
  if (v1.size() != v2.size())
  {
    assert(false);
    return;
  }
  for (int i = 0; i < v1.size(); i++)
  {
    Assert(v1[i], v2[i]);
  }
}
int main()
{
   vector<int> nums;
   int k, maxChanges;
   {
     nums = { 1,0,1,0,1 }, k = 3, maxChanges = 0;
     auto res = Solution().minimumMoves(nums, k, maxChanges);
     Assert(4, res);
   }
  {
     nums = { 1, 1, 0, 0, 0, 1, 1, 0, 0, 1 }, k = 3, maxChanges = 1;
     auto res = Solution().minimumMoves(nums, k, maxChanges);
    Assert(3, res);
  }
  {
    nums = { 0, 0, 0, 0 }, k = 2, maxChanges = 3;
    auto res = Solution().minimumMoves(nums, k, maxChanges);
    Assert(4, res);
  }
}

优化

如果有情况四，情况一二三，可以统一。

class Solution {
public:
  long long minimumMoves(const vector<int>& nums, const int k, const int maxChanges) {
    m_c = nums.size();
    queue<int> que1Index;
    for (int i = 0; i < m_c; i++)
    {
      if (nums[i])
      {
        que1Index.emplace(i);
      }
    }
    vector<long long> vPre(nums.size());//距离 nums[i]最近的iNum个1的距离
    const int i1Num = k - maxChanges;   
    queue<int> que1, que2;
    while(que1Index.size() && (que2.size() < i1Num))
    {//计算vPre[0]
      que2.emplace(que1Index.front());
      vPre[0] += que1Index.front();
      que1Index.pop();
    }
    for (int i2 = 1; i2 < m_c; i2++)
    {
      vPre[i2] = vPre[i2 - 1] + que1.size() - que2.size();
      if (que2.size() && (que2.front() < i2))
      {//队列二向队列一移动
        que1.emplace(que2.front());
        que2.pop();
        vPre[i2] += 1 + 1;
      }
      while (que1.size() && que1Index.size() && (i2 - que1.front() > que1Index.front() - i2))
      {
        vPre[i2] -= (i2 - que1.front());
        vPre[i2] += que1Index.front() - i2;
        que2.emplace(que1Index.front());
        que1.pop();
        que1Index.pop();
      }
    }
    long long llRet = LLONG_MAX;;
    for (int i = 0; i < nums.size(); i++)
    {
      int iNeed = k - nums[i];  
      long long llDo = 0;
      if ((iNeed > 0) && (i > 0) && nums[i - 1])
      {
        iNeed--; llDo++;
      }
      if ((iNeed > 0) && (i+1 < nums.size()) && nums[i + 1])
      {
        iNeed--; llDo++;
      }
      if (maxChanges >= iNeed)
      {
        llDo += 2 * iNeed;
      }
      else
      {
        llDo = vPre[i] + maxChanges * 2;
      }
      llRet = min(llRet, llDo);
    }
    return llRet;
  }
  int m_c;
};

中位数贪心

如果某个数距离m个数距离最短，那么它一定是正中间的数。也就是情况四：左边m/2个1，右边m-1-m/2个1。用前缀和计算。

nums[index]不会为0。

一，000。根据中位数贪心，相比换到这m个数的中心。情况一二四一定不会有优势。情况一二也没优先。

二，100，同上。

三，101。假定101左边有m1个1，右边有m2个1。如果m1 <= m2，移到101 ，否则移到101

除101外的数距离全部变短。101的距离从:1,1变成 0,2 持平。

如何计算左边m1个和

vPre[i]记录前i个1到0的距离。

故第j个1前面的m1个。m1 × \times× index[j]- (vPre[j]-vPre[j-m1])

故第j个1后面的m2个 vPre[j+1+m2] - vPre[j+1] - m2× \times×index[j]

特例

全部是0。

代码

class Solution {
public:
  long long minimumMoves(const vector<int>& nums, const int k, const int maxChanges) {
    m_c = nums.size();
    vector<int> v1Index;
    for (int i = 0; i < m_c; i++)
    {
      if (nums[i])
      {
        v1Index.emplace_back(i);
      }
    }
    if (v1Index.empty())
    {
      return k * 2;
    }
    vector<long long> vPre = { 0 };
    for (const auto& n : v1Index)
    {
      vPre.emplace_back(vPre.back() + n);
    }
    long long llRet = LLONG_MAX;;
    for (int j = 0 ; j < v1Index.size();j++ )
    {
      const int& i = v1Index[j];
      int iNeed = k - nums[i];  
      long long llDo = 0;
      if ((iNeed > 0) && (i > 0) && nums[i - 1])
      {
        iNeed--; llDo++;
      }
      if ((iNeed > 0) && (i+1 < nums.size()) && nums[i + 1])
      {
        iNeed--; llDo++;
      }
      if (maxChanges >= iNeed)
      {
        llDo += 2 * iNeed;
      }
      else
      {
        const int m = k - maxChanges;
        const long long m1 = m / 2;
        const long long m2 = m - m1 - 1;
        if ((j - m1 < 0) || (j + 1 + m2 >= vPre.size()))
        {
          continue;
        }
        const long long llLeft = m1 * i - (vPre[j] - vPre[j - m1]);
        const long long llRight = vPre[j + 1 + m2] - vPre[j + 1] - m2 * i;
        llDo = llLeft + llRight + maxChanges * 2;
      }
      llRet = min(llRet, llDo);
    }
    return llRet;
  }
  int m_c;
};

扩展阅读

视频课程

有效学习：明确的目标及时的反馈拉伸区（难度合适），可以先学简单的课程，请移步CSDN学院，听白银讲师（也就是鄙人）的讲解。

https://edu.csdn.net/course/detail/38771

如何你想快速形成战斗了，为老板分忧，请学习C#入职培训、C++入职培训等课程

https://edu.csdn.net/lecturer/6176

我想对大家说的话
闻缺陷则喜是一个美好的愿望，早发现问题，早修改问题，给老板节约钱。
子墨子言之：事无终始，无务多业。也就是我们常说的专业的人做专业的事。
如果程序是一条龙，那算法就是他的是睛

测试环境

操作系统：win7 开发环境： VS2019 C++17

或者操作系统：win10 开发环境： VS2022 C++17

如无特殊说明，本算法用**C++**实现。

【对顶队列】【中位数贪心】【前缀和】3086. 拾起 K 个 1 需要的最少行动次数