Minimum Adjacent Swaps for K Consecutive Ones LeetCode solution in Python

1
class Solution:
2
    def minMoves(self, nums: List[int], k: int) -> int:
3
        p = [i for i, v in enumerate(nums) if v == 1]
4
        # p[i]: the position of i-th 1
5
        n = len(p)
6
        presum = [0] * (n + 1)
7
        # presum[i]: sum(p[0]...p[i-1])
8
        for i in range(n):
9
            presum[i + 1] = presum[i] + p[i]
10

11
        res = inf
12

13
        # sliding window
14
        if k % 2 == 1:
15
            # if odd
16
            radius = (k - 1) // 2
17
            for i in range(radius, n - radius):
18
                # i-radius ... i ... i+radius
19
                # move radius to i
20
                # i+1, ..., i+radius
21
                right = presum[i + radius + 1] - presum[i + 1]
22
                # i-radius, ..., i-1
23
                left = presum[i] - presum[i - radius]
24
                res = min(res, right - left)
25
            return res - radius * (radius + 1)
26
        else:
27
            # even
28
            radius = (k - 2) // 2
29
            for i in range(radius, n - radius - 1):
30
                # i-radius ... i i+1 ... i+radius+1
31
                # move radius to i (moving to i+1 is also OK)
32
                # i+1, ..., i+radius+1
33
                right = presum[i + radius + 2] - presum[i + 1]
34
                # i-radius, ..., i-1
35
                left = presum[i] - presum[i - radius]
36
                res = min(res, right - left - p[i])
37
            return res - radius * (radius + 1) - (radius + 1)