小L的元素修改问题 | 豆包MarsCode AI刷题

问题描述

小R拿到了一个数组，她可以进行如下操作：使得一个元素加1，另一个元素减1。她希望最终数组的每个元素大小都在 [l, r] 的范围内。小R想知道，最少需要多少次操作可以达到目标。

如果无法通过有限次操作使所有元素都落在指定范围内，则返回 -1。

---

解题思路

• 判断可行性：

  • 所有元素的总和必须在区间 $[n \times l, n \times r]$ 内。如果总和不满足此条件，则说明无论如何操作也无法使数组元素满足要求。

• 最少操作思路：

  • 每次操作将数组中的最大值向下调整-1，最小值向上调整 +1，同时计数，逐步使得所有元素的值落入区间 $[l, r]$。
  • 为了高效找到当前的最大值和最小值，使用 最大堆 和 最小堆 分别维护数组中的最大值和最小值（堆只能从顶端取数据）。

---

算法设计

• 初始化：

  • 计算数组总和 sum。
  • 如果 sum 不在区间 $[n \times l, n \times r]$ 内，直接返回 -1。
  • 初始化最大堆和最小堆，存储数组中的元素。

• 调整过程：

  • 不合法情况已在之前排除，此时均为合法情况。

  • 获取堆顶元素作为当前的最大值 maxVal 和最小值 minVal。

  • 判断：如果 maxVal 和 minVal 均在区间 $[l, r]$ 内，停止操作，返回操作次数。

  • 更新：

    • maxVal --
    • minVal ++

  • 更新最大堆和最小堆，记录操作次数。

---

代码实现

#include <iostream>
#include <vector>
#include <queue>
#include <numeric>

int solution(int n, int l, int r, std::vector<int> a) {
    // 计算数组总和
    int sum = std::accumulate(a.begin(), a.end(), 0);

    // 检查是否可行
    if (sum < n * l || sum > n * r)
        return -1;

    // 使用优先队列处理最大值和最小值
    std::priority_queue<int> maxHeap; // 最大堆，用于获取最大值
    std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; // 最小堆，用于获取最小值

    // 初始化堆
    for (int num : a)
    {
        maxHeap.push(num);
        minHeap.push(num);
    }

    int operations = 0;

    // 开始调整
    while (true)
    {
        int maxVal = maxHeap.top(); // 获取当前最大值
        int minVal = minHeap.top(); // 获取当前最小值

        // 检查是否所有元素都在范围内
        if (minVal >= l && maxVal <= r)
            return operations;

        // 移除堆顶元素
        maxHeap.pop();
        minHeap.pop();

        // 调整最大值和最小值
        maxVal--; // 最大值减 1
        minVal++; // 最小值加 1

        // 更新堆
        maxHeap.push(maxVal);
        minHeap.push(minVal);

        operations++;
    }
}

int main() {
    std::cout << (solution(2, 3, 5, {1, 2}) == -1) << std::endl;
    std::cout << (solution(3, 4, 6, {3, 6, 5}) == 1) << std::endl;
    std::cout << (solution(4, 2, 8, {1, 10, 2, 6}) == 2) << std::endl;
}

---

复杂度分析

1. 时间复杂度：

   • 初始化堆的复杂度为 $O(n \log n)$。
   • 每次调整操作的复杂度为 $O(\log n)$，总调整次数最多为 $O(n)$，因此总时间复杂度为 $O(n \log n)$。

2. 空间复杂度：

   • 额外使用了两个堆，每个堆的空间复杂度为 $O(n)$。

---

示例分析

示例 1：

输入：n = 2, l = 3, r = 5, a = {1, 2}

• 初始总和：$1 + 2 = 3$，小于 $2 \times 3 = 6$，直接返回 $-1$。

输出：-1

示例 2：

输入：n = 3, l = 4, r = 6, a = {3, 6, 5}

• 初始总和：$3 + 6 + 5 = 14$，满足 $3 \times 4 \leq 14 \leq 3 \times 6$。

• 第一次操作：

  • 最大值 $6$ 减 1，最小值 $3$ 加 1。
  • 数组变为 $[4, 5, 5]$。

• 此时所有元素均在区间内，操作完成。

输出：1

示例 3：

输入：n = 4, l = 2, r = 8, a = {1, 10, 2, 6}

• 初始总和：$1 + 10 + 2 + 6 = 19$，满足 $4 \times 2 \leq 19 \leq 4 \times 8$。

• 第一次操作：

  • 最大值 $10$ 减 1，最小值 $1$ 加 1。
  • 数组变为 $[2, 9, 2, 6]$。

• 第二次操作：

  • 最大值 $9$ 减 1，最小值 $2$ 加 1。
  • 数组变为 $[3, 8, 2, 6]$。

• 此时所有元素均在区间内，操作完成。

输出：2

---

总结

• 本题的关键在于判断总和是否满足条件，以及通过优先队列高效调整最大值和最小值。
• 使用堆结构优化后，解决方案的效率显著提升，适用于大规模输入。