https://leetcode.com/problems/minimum-adjacent-swaps-for-k-consecutive-ones/ from Google

Thinking process:

First look at the constraints. We shouldn’t use any algorithm slower than O(NlogN). This is kind of hard at the beginning. Maybe we can find a solution of a simpler question first? Let’s say if we have 5 1s in the array, and we want to find the minimum moves to make them 5 consecutive 1’s. What’s the solution?

The a,b,c,d above means how many zeros are there between two ones, as a result, they all >= 0. From the equations, we can know that merging to index 3 is best, which means that the minimum moves always happen when we merge all ones to the middle one.

After knowing this, it becomes easier for this question as we know how to do the merge. The only thing we don’t know is which set of the k ones is the solution. We can check this by using a slide window. So, again, if k = 5, and we’re at a window with a,b,c,d zeros in between, we want to slide the window to right so we remove a and add e, just like below.

You can think (a+b)+b as lall, and a+b as the l, which is the sum of all zeros of the left part. Similarly, c+(c+d) is rall, and c+d is r. Then, with some math like this, we can figure out all possible solution in O(N) time.

Here is the code.

 1class Solution {
 2public:
 3    int minMoves(vector<int>& nums, int k) {
 4        int N = nums.size();
 5        int pre = -1;
 6        vector<int> ar;
 7        for (int i=0; i<N; i++) {
 8            if (nums[i] == 1) {
 9                if (pre == -1) pre = i;
10                else {
11                    ar.push_back(i-pre-1);
12                    pre = i;
13                }
14            }
15        }
16        if (k == 1) return 0;
17        int ans;
18        int L = ar.size();
19        int l = 0, r = 0, mid = k/2, lall = 0, rall = 0;
20        for (int i=(k/2)-1; i>=0; i--) {
21            l += ar[i];
22            lall += l;
23        }
24        for (int i=k/2; i<k-1; i++) {
25            r += ar[i];
26            rall += r;
27        }
28        ans = lall + rall;
29        for (int i=k-1; i<L; i++, mid++) {
30            lall = lall - l + ar[mid]*(k/2);
31            l = l - ar[i-k+1] + ar[mid];
32            r = r - ar[mid] + ar[i];
33            rall = rall - ar[mid]*(k/2 - ((k%2) == 0)) + r;
34            ans = min(ans, lall+rall);
35        }
36        return ans;
37    }
38};