3107. 使数组中位数等于 K 的最少操作数

题目描述

给你一个整数数组 nums 和一个 非负 整数 k 。一次操作中，你可以选择任一元素 加 1 或者减 1 。

请你返回将 nums 中位数 变为 k 所需要的 最少 操作次数。

一个数组的中位数指的是数组按非递减顺序排序后最中间的元素。如果数组长度为偶数，我们选择中间两个数的较大值为中位数。

 

示例 1：

输入：nums = [2,5,6,8,5], k = 4

输出：2

解释：我们将 nums[1] 和 nums[4] 减 1 得到 [2, 4, 6, 8, 4] 。现在数组的中位数等于 k 。

示例 2：

输入：nums = [2,5,6,8,5], k = 7

输出：3

解释：我们将 nums[1] 增加 1 两次，并且将 nums[2] 增加 1 一次，得到 [2, 7, 7, 8, 5] 。

示例 3：

输入：nums = [1,2,3,4,5,6], k = 4

输出：0

解释：数组中位数已经等于 k 了。

 

提示：

• 1 <= nums.length <= 2 * 105
• 1 <= nums[i] <= 109
• 1 <= k <= 109

解法

方法一：贪心 + 排序

我们首先对数组 \(nums\) 进行排序，然后找到中位数的位置 \(m\)，那么初始时我们需要的操作次数就是 \(|nums[m] - k|\)。

接下来，我们分情况讨论：

• 如果 \(nums[m] \gt k\)，那么 \(m\) 右侧的元素都大于等于 \(k\)，我们只需要将 \(m\) 左侧的元素中，大于 \(k\) 的元素减小到 \(k\) 即可。
• 如果 \(nums[m] \le k\)，那么 \(m\) 左侧的元素都小于等于 \(k\)，我们只需要将 \(m\) 右侧的元素中，小于 \(k\) 的元素增大到 \(k\) 即可。

时间复杂度 \(O(n \times \log n)\)，空间复杂度 \(O(\log n)\)。其中 \(n\) 是数组 \(nums\) 的长度。

Python3JavaC++GoTypeScript

class Solution:
    def minOperationsToMakeMedianK(self, nums: List[int], k: int) -> int:
        nums.sort()
        n = len(nums)
        m = n >> 1
        ans = abs(nums[m] - k)
        if nums[m] > k:
            for i in range(m - 1, -1, -1):
                if nums[i] <= k:
                    break
                ans += nums[i] - k
        else:
            for i in range(m + 1, n):
                if nums[i] >= k:
                    break
                ans += k - nums[i]
        return ans

class Solution {
    public long minOperationsToMakeMedianK(int[] nums, int k) {
        Arrays.sort(nums);
        int n = nums.length;
        int m = n >> 1;
        long ans = Math.abs(nums[m] - k);
        if (nums[m] > k) {
            for (int i = m - 1; i >= 0 && nums[i] > k; --i) {
                ans += nums[i] - k;
            }
        } else {
            for (int i = m + 1; i < n && nums[i] < k; ++i) {
                ans += k - nums[i];
            }
        }
        return ans;
    }
}

class Solution {
public:
    long long minOperationsToMakeMedianK(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());
        int n = nums.size();
        int m = n >> 1;
        long long ans = abs(nums[m] - k);
        if (nums[m] > k) {
            for (int i = m - 1; i >= 0 && nums[i] > k; --i) {
                ans += nums[i] - k;
            }
        } else {
            for (int i = m + 1; i < n && nums[i] < k; ++i) {
                ans += k - nums[i];
            }
        }
        return ans;
    }
};

func minOperationsToMakeMedianK(nums []int, k int) (ans int64) {
    sort.Ints(nums)
    n := len(nums)
    m := n >> 1
    ans = int64(abs(nums[m] - k))
    if nums[m] > k {
        for i := m - 1; i >= 0 && nums[i] > k; i-- {
            ans += int64(nums[i] - k)
        }
    } else {
        for i := m + 1; i < n && nums[i] < k; i++ {
            ans += int64(k - nums[i])
        }
    }
    return
}

func abs(x int) int {
    if x < 0 {
        return -x
    }
    return x
}

function minOperationsToMakeMedianK(nums: number[], k: number): number {
    nums.sort((a, b) => a - b);
    const n = nums.length;
    const m = n >> 1;
    let ans = Math.abs(nums[m] - k);
    if (nums[m] > k) {
        for (let i = m - 1; i >= 0 && nums[i] > k; --i) {
            ans += nums[i] - k;
        }
    } else {
        for (let i = m + 1; i < n && nums[i] < k; ++i) {
            ans += k - nums[i];
        }
    }
    return ans;
}

评论