BinarySearch 0548 Group the Ones

Problem statement

https://binarysearch.com/problems/Group-the-Ones/

Solution

Equal to 1703. Minimum Adjacent Swaps for K Consecutive Ones.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(self, A, k):
        A = [i for i, a in enumerate(A) if a == "1"]
        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)