LeetCode每日一题(1703. Minimum Adjacent Swaps for K Consecutive Ones)

You are given an integer array, nums, and an integer k. nums comprises of only 0’s and 1’s. In one move, you can choose two adjacent indices and swap their values.

Return the minimum number of moves required so that nums has k consecutive 1’s.

Example 1:

Input: nums = [1,0,0,1,0,1], k = 2
Output: 1

Explanation: In 1 move, nums could be [1,0,0,0,1,1] and have 2 consecutive 1’s.

Example 2:

Input: nums = [1,0,0,0,0,0,1,1], k = 3
Output: 5

Explanation: In 5 moves, the leftmost 1 can be shifted right until nums = [0,0,0,0,0,1,1,1].

Example 3:

Input: nums = [1,1,0,1], k = 2
Output: 0

Explanation: nums already has 2 consecutive 1’s.

Constraints:

• 1 <= nums.length <= 105
• nums[i] is 0 or 1.
• 1 <= k <= sum(nums)

---

用例子来说吧

nums = [1, 0, 0, 1, 0, 1, 0, 1, 1, 0]
k = 3

我们先把 nums 中所有 1 的 index 都收集起来得到:

one_positions = [0, 3, 5, 7, 8]

这时候我们用一个长度等于 k 的 sliding window 来遍历 one_positions, 会得到 3 个 window， [0, 3, 5], [3, 5, 7]和[5, 7, 8]

然后我们就可以开始计算移动步数了

[0, 3, 5] 最终移动成 [2, 3, 4] 一共用了 3 步
[3, 5, 7] 最终移动成 [4, 5, 6] 一共用了 2 步
[5, 7, 8] 最终移动成 [6, 7, 8] 一共用了 1 步

这里有一个非常重要的一点， 就是最小步数一定是两侧的数字往中间点靠拢得到的

这样， 我们的问题也就变成了， 两侧的数字 1 往中间的数字 1 移动形成连续数字 1，需要多少步

看 nums 和 k 的取值范围， 我们应该清楚， 每个 window 遍历一次进行计算的话肯定会超时， 所以我们需要想其他办法。

两个元素之间的距离很容易算，比如[0, 3, 5]中 0 到 3 的距离就是 3 - 0 = 3, 也就是将 0 移动到 3 的位置需要 3 步， 但是我们不需要将 0 移动到 3， 我们只需要将 0 移动到 2 就可以了。这时我们假设我们需要将[0, 3, 5]移动成[3, 4, 5], 我们来看一下元素位置和目标位置以及移动步数的关系。

首先 3 要移动到 4， 这样是 1 步， 然后 0 要移动到 3， 这样是 3 步， 这样我们可以看出， 从右向左来看， 每个要移动的元素它们的目标位置是依次减 1 的。但是这样我们还是需要遍历来进行计算， 效率依然不高。 我们可以换一种算法。

我们先算每个元素移动到 5 的步数

0 到 5 => 5
3 到 5 => 2

此时我们肯定是移动多了， 具体移动多了多少呢， 我们看最终结果的差异, 本来我们的目标是[3, 4, 5], 现在移动得到了[5, 5, 5], 实际多移动的就是(5 - 3) + (5 - 4) = 2 + 1, 以此类推， 如果我们将原始数组换成[0, 1, 3, 5]的话， 目标数组应该是[2, 3, 4, 5], 但是我们都移动到 5 的话就是[5, 5, 5, 5], 这样我们就多移动了(5 - 2) + (5 - 3) + (5 - 4) = 3 + 2 + 1, 如果是 5 个数的话就会变成 4 + 3 + 2 + 1, 看着眼熟吧？就是 n * (n - 1) / 2, 其中 n 是数组长度。

这样我们就可以用 sum(steps_to_right_most) - n * (n - 1) / 2 来算实际需要的移动步数

那 sum(steps_to_right_most)怎么算呢？这时候我们可以用 prefix sum 的方法。比较容易想， 就不展开说了。

问题基本解决， 但是我们还是要考虑另一个问题， 就是我们上面的例子中 k 是奇数， 如果 k 是偶数， 我们会有两个中间点， 我们将左右两侧的数字分别移动到这两个中间点的两侧,然后再加上 k / 2 * 两个中间点移动到一起所需要的步数就是最终答案

---

impl Solution {
    pub fn min_moves(nums: Vec<i32>, k: i32) -> i32 {
        let one_positions: Vec<usize> = nums.iter().enumerate().filter(|(_, &v)| v == 1).map(|(i, _)| i).collect();
        let mut prefix: Vec<usize> = one_positions
            .iter()
            .scan(0, |s, v| {
                *s += *v;
                Some(*s)
            })
            .collect();
        prefix.insert(0, 0);
        let mut ans = usize::MAX;
        if k % 2 == 0 {
            for i in 0..one_positions.len() - k as usize + 1 {
                let l = i;
                let r = i + k as usize - 1;
                let lm = (l + r) / 2;
                let rm = lm + 1;
                let left_steps = one_positions[lm] * (lm - l + 1) - (prefix[lm + 1] - prefix[l]) - (lm - l + 1) * (lm - l) / 2;
                let right_steps = prefix[r + 1] - prefix[rm] - one_positions[rm] * (r - rm + 1) - (r - rm + 1) * (r - rm) / 2;
                ans = ans.min(left_steps + right_steps + (one_positions[rm] - one_positions[lm] - 1) * k as usize / 2);
            }
            return ans as i32;
        }
        for i in 0..one_positions.len() - k as usize + 1 {
            let l = i;
            let r = i + k as usize - 1;
            let m = (l + r) / 2;
            let left_steps = one_positions[m] * (m - l + 1) - (prefix[m + 1] - prefix[l]) - (m - l + 1) * (m - l) / 2;
            let right_steps = prefix[r + 1] - prefix[m] - one_positions[m] * (r - m + 1) - (r - m + 1) * (r - m) / 2;
            ans = ans.min(left_steps + right_steps);
        }
        ans as i32
    }
}