Learn AI Create the Future — with Pranay
Main site

LeetCode solutions

3327. Minimum Moves to Pick K Ones

My accepted Python solution to LeetCode problem 3327, Minimum Moves to Pick K Ones, running in 107ms.

  • Difficulty: Hard
  • Python
  • Runtime 107ms
  • Memory 27.9MB

Read the problem on LeetCode View on GitHub

The problem statement is LeetCode’s and stays on their site. What follows is my accepted solution.

Python

Accepted on LeetCode — runtime 107ms, memory 27.9MB, accepted 2026-01-02.

python
class Solution:
    def minimumMoves(self, nums: List[int], k: int, maxChanges: int) -> int:
        n = len(nums)
        ones = []
        
        # Find all positions of 1s
        for i in range(n):
            if nums[i] == 1:
                ones.append(i)
        
        # Count max consecutive 1s (at most 3 that we can pick for free or with 1 move)
        max_consec = 0
        if ones:
            consec = 1
            for i in range(1, len(ones)):
                if ones[i] == ones[i-1] + 1:
                    consec += 1
                    max_consec = max(max_consec, consec)
                else:
                    consec = 1
            max_consec = max(max_consec, 1)
        
        # Special case: if we can pick all k using changes and at most 3 from consecutive
        # From immediate neighbors (or standing position), we can pick up to 3 ones (1 from standing + 2 adjacent)
        # Each one from adjacent costs 1 move
        # Each one from maxChanges costs 2 moves (1 to create, 1 to pick)
        
        # If k <= 3 and we have enough consecutive or can use changes
        c = min(3, max_consec, k)  # ones we can get from consecutive segment
        
        if k <= c + maxChanges:
            # We can satisfy k using c from consecutive and (k-c) from changes
            # Getting c ones costs c-1 moves (first is free)
            # Getting remaining k-c from changes costs 2*(k-c) moves
            return max(0, c - 1) + 2 * max(0, k - c)
        
        # Otherwise, we need to use sliding window on the positions of 1s
        # We need (k - maxChanges) ones from existing positions
        need = k - maxChanges
        
        # Build prefix sum of positions
        m = len(ones)
        prefix = [0] * (m + 1)
        for i in range(m):
            prefix[i + 1] = prefix[i] + ones[i]
        
        # Sliding window of size 'need', find minimum cost to collect them
        # Cost of moving ones to median position
        result = float('inf')
        
        for i in range(m - need + 1):
            # Window is ones[i:i+need]
            # Median is at position i + (need-1)//2
            mid = i + (need - 1) // 2
            median_pos = ones[mid]
            
            # Cost to move left half to median
            left_count = mid - i
            left_sum = prefix[mid] - prefix[i]
            left_cost = left_count * median_pos - left_sum
            
            # Cost to move right half to median
            right_count = (i + need - 1) - mid
            right_sum = prefix[i + need] - prefix[mid + 1]
            right_cost = right_sum - right_count * median_pos
            
            # Total cost for existing ones + 2*maxChanges for created ones
            total = left_cost + right_cost + 2 * maxChanges
            
            # We may be able to get some adjacent ones for cheaper
            # If median has adjacent 1s, they cost 1 instead of distance
            result = min(result, total)
        
        return result

Source