Leetcode 1703. Minimum Adjacent Swaps for K Consecutive Ones

Problem statement

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

Solution

The idea is the following: let us look at all places where we have 1. Then we need to find the window of size k, such that sum between elements in this group and its median is the smallest. For one window we can check it in `O(1)

Complexity

It is O(n) for time and space.

Code

class Solution:
    def minMoves(self, A, k):
        A = [i for i, a in enumerate(A) if a]
        ans = float("inf")
        acc = [0] + list(accumulate(A))
     
        for i in range(len(A) - k + 1):
            ans = min(ans, acc[i + k] - acc[k//2 + i] - acc[(k + 1)//2 + i] + acc[i])
        return ans - (k//2) * ((k + 1)//2)