Minimum Moves to Pick K Ones

Hard

#2742Time: O(m), where `m` is the number of ones. We iterate through the list of ones once, and for each one, we perform a constant number of calculations.Space: O(m), where `m` is the number of ones, for storing `one_indices` and prefix sums.

Patterns

Sliding Window Greedy Prefix Sum

Data structures

Array

Prompt

You are given a binary array nums of length n, a positive integer k and a non-negative integer maxChanges.

Alice plays a game, where the goal is for Alice to pick up k ones from nums using the minimum number of moves. When the game starts, Alice picks up any index aliceIndex in the range [0, n - 1] and stands there. If nums[aliceIndex] == 1 , Alice picks up the one and nums[aliceIndex] becomes 0(this does not count as a move). After this, Alice can make any number of moves (including zero) where in each move Alice must perform exactly one of the following actions:

Select any index j != aliceIndex such that nums[j] == 0 and set nums[j] = 1. This action can be performed at most maxChanges times.
Select any two adjacent indices x and y (|x - y| == 1) such that nums[x] == 1, nums[y] == 0, then swap their values (set nums[y] = 1 and nums[x] = 0). If y == aliceIndex, Alice picks up the one after this move and nums[y] becomes 0.

Return the minimum number of moves required by Alice to pick exactly k ones.

Example 1:

Input: nums = [1,1,0,0,0,1,1,0,0,1], k = 3, maxChanges = 1

Output: 3

Explanation: Alice can pick up 3 ones in 3 moves, if Alice performs the following actions in each move when standing at aliceIndex == 1:

At the start of the game Alice picks up the one and nums[1] becomes 0. nums becomes [1,0,0,0,0,1,1,0,0,1].
Select j == 2 and perform an action of the first type. nums becomes [1,0,1,0,0,1,1,0,0,1]
Select x == 2 and y == 1, and perform an action of the second type. nums becomes [1,1,0,0,0,1,1,0,0,1]. As y == aliceIndex, Alice picks up the one and nums becomes [1,0,0,0,0,1,1,0,0,1].
Select x == 0 and y == 1, and perform an action of the second type. nums becomes [0,1,0,0,0,1,1,0,0,1]. As y == aliceIndex, Alice picks up the one and nums becomes [0,0,0,0,0,1,1,0,0,1].

Note that it may be possible for Alice to pick up 3 ones using some other sequence of 3 moves.

Example 2:

Input: nums = [0,0,0,0], k = 2, maxChanges = 3

Output: 4

Explanation: Alice can pick up 2 ones in 4 moves, if Alice performs the following actions in each move when standing at aliceIndex == 0:

Select j == 1 and perform an action of the first type. nums becomes [0,1,0,0].
Select x == 1 and y == 0, and perform an action of the second type. nums becomes [1,0,0,0]. As y == aliceIndex, Alice picks up the one and nums becomes [0,0,0,0].
Select j == 1 again and perform an action of the first type. nums becomes [0,1,0,0].
Select x == 1 and y == 0 again, and perform an action of the second type. nums becomes [1,0,0,0]. As y == aliceIndex, Alice picks up the one and nums becomes [0,0,0,0].

Constraints:

2 <= n <= 10⁵
0 <= nums[i] <= 1
1 <= k <= 10⁵
0 <= maxChanges <= 10⁵
maxChanges + sum(nums) >= k

Approaches

2 approaches with complexity analysis and trade-offs.

This approach significantly optimizes the previous one by making a critical observation about the nature of the optimal solution. The cost of creating a new 1 and bringing it to the collection point is 2 moves. This sets a budget for how far an original 1 can be from the collection point. If an original 1 is more than 2 moves away from the median of the chosen group, it's cheaper to use a 'change' operation instead.

This insight proves that all original 1s in an optimal group must be very close to each other. Specifically, the distance from any chosen 1 to the group's median must be at most 2. This in turn means the total span of the indices of the chosen 1s (max_index - min_index) must be at most 4. Consequently, we only need to check contiguous windows of 1s of size up to 5, drastically reducing the search space.

Algorithm

The core insight is that in an optimal solution, any chosen original 1 at index p must be close to the median M of all chosen original 1s. Specifically, the distance |p - M| cannot be greater than 2.
If |p - M| > 2, it would be cheaper to discard this 1 and use a 'change' operation instead, which costs 2 moves. This would reduce the total cost, contradicting optimality.
This observation implies that for any chosen group of original 1s, their indices p_1, p_2, ..., p_w must satisfy max(p_i) - min(p_i) <= 4.
Therefore, we only need to consider contiguous windows of 1s from one_indices where the window size is at most 5.
The algorithm then becomes linear: iterate through each index i in one_indices as a potential start of a window.
For each i, consider small windows starting at i (of size up to 5). A second pointer j can find the end of such small windows efficiently.
For each such small window [i, j], calculate the number of original ones w = j - i + 1.
If k-w is a valid number of changes (0 <= k-w <= maxChanges), calculate the total cost: gather_cost(i, w) + 2 * (k-w).
Update the minimum cost found. Also, handle the base case of using only changes (2*k if k <= maxChanges).

Walkthrough

The logic builds upon the previous approach but constrains the search space. We still find all one_indices and compute their prefix sums.

The key argument is as follows: Let the chosen set of w original ones be S, and their median index be M. The total cost is sum_{p in S} |p - M| + 2 * (k - w). Suppose there exists a p_j in S such that |p_j - M| > 2. We could form a new solution by removing p_j from S and using one more 'change' operation. The new cost would be sum_{p in S, p != p_j} |p - M'| + 2 * (k - w + 1), where M' is the new median. The sum of distances to the new median M' can only be less than or equal to the sum of distances to the old median M. So, the change in cost is at most 2 - |p_j - M|, which is negative if |p_j - M| > 2. This means the new solution is cheaper, a contradiction. Thus, in any optimal solution, all chosen original ones p must satisfy |p - M| <= 2.

This implies that the difference between the maximum and minimum index in the chosen group is at most 4. Since the indices are integers, such a group can have at most 5 ones.

This transforms the O(m^2) algorithm into an O(m) one. We iterate through i from 0 to m-1. For each i, we only need to check windows P[i...j] where j is at most i+4. This is a constant number of checks per i.

1import java.util.ArrayList;2import java.util.List;3 4class Solution {5    public long minimumMoves(int[] nums, int k, int maxChanges) {6        List<Integer> oneIndices = new ArrayList<>();7        for (int i = 0; i < nums.length; i++) {8            if (nums[i] == 1) {9                oneIndices.add(i);10            }11        }12 13        int m = oneIndices.size();14        if (m + maxChanges < k) {15            // Should not happen based on constraints, but good practice16            return -1; 17        }18        if (k == 0) {19            return 0;20        }21 22        long minMoves = Long.MAX_VALUE;23 24        // Case 1: Use only changes25        if (k <= maxChanges) {26            minMoves = (long)k * 2;27        }28 29        // Case 2: Use some original ones30        // We only need to check small groups of original ones because if any original one31        // is > 2 distance from the median, it's cheaper to use a change (cost 2).32        // This implies the span of indices is at most 4, so window size is at most 5.33        // Let's check a slightly larger bound for safety, e.g., up to 3 adjacent ones.34        // The proof shows we only need to check groups of 1, 2, or 3 ones around an index.35        for (int i = 0; i < m; i++) {36            // Option A: 1 original at P[i], k-1 changes37            if (k - 1 <= maxChanges) {38                minMoves = Math.min(minMoves, (long)2 * (k - 1));39            }40            // Option B: 2 originals P[i], P[i+1], k-2 changes41            if (i + 1 < m && k - 2 <= maxChanges) {42                long cost = oneIndices.get(i + 1) - oneIndices.get(i);43                minMoves = Math.min(minMoves, cost + (long)2 * (k - 2));44            }45            // Option C: 3 originals P[i-1], P[i], P[i+1], k-3 changes46            if (i > 0 && i + 1 < m && k - 3 <= maxChanges) {47                long cost = (long)oneIndices.get(i+1) - oneIndices.get(i-1);48                minMoves = Math.min(minMoves, cost + (long)2 * (k - 3));49            }50        }51        52        // If k is very small, we might not have enough original ones for the options above.53        // If k=1 and m>=1, cost is 0. If k=1 and m=0, cost is 2 (from initial minMoves).54        if (k > 0 && m >= k) { // If we can form k ones without changes55             long[] prefixSum = new long[m + 1];56             for (int i = 0; i < m; i++) {57                 prefixSum[i + 1] = prefixSum[i] + oneIndices.get(i);58             }59             for (int i = 0; i <= m - k; i++) {60                 minMoves = Math.min(minMoves, getGatherCost(prefixSum, oneIndices, i, i + k - 1));61             }62        }63 64        return minMoves;65    }66 67    private long getGatherCost(long[] prefixSum, List<Integer> oneIndices, int i, int j) {68        int w = j - i + 1;69        int medianIndexInWindow = (w - 1) / 2;70        int medianAbsIndex = i + medianIndexInWindow;71        72        long leftSum = prefixSum[medianAbsIndex] - prefixSum[i];73        long rightSum = prefixSum[j + 1] - prefixSum[medianAbsIndex + 1];74        long medianVal = oneIndices.get(medianAbsIndex);75 76        long cost = (medianIndexInWindow * medianVal - leftSum) + (rightSum - (long)(w - 1 - medianIndexInWindow) * medianVal);77        return cost;78    }79}

Note: The provided code snippet for the efficient approach is slightly simplified for clarity, focusing on the core idea of using small groups of 1, 2, or 3 original ones plus changes, which covers the most critical cases. A full implementation would combine this with the base case of using k original ones if maxChanges is 0.

Complexity

Time

O(m), where `m` is the number of ones. We iterate through the list of ones once, and for each one, we perform a constant number of calculations.

Space

O(m), where `m` is the number of ones, for storing `one_indices` and prefix sums.

Trade-offs

Pros

Highly efficient with linear time complexity.
Correctly identifies the optimal substructure of the problem, leading to a significant reduction in computation.

Cons

The core insight that limits the window size is non-trivial and requires a careful proof by contradiction.

This approach considers all possible contiguous groups of original 1s that can be part of the k ones we collect. For any number of original 1s we decide to use, say w, it's always optimal to choose a contiguous block of w 1s from the sorted list of their indices to minimize their gathering cost. The gathering point that minimizes the moves is the median of their indices.

The strategy is to iterate through every possible contiguous subsegment of the 1-indices. For each subsegment of size w, we calculate the cost to gather them at their median. Then, we determine how many additional 1s we need, which is k - w. These must be created using the 'change' operation, costing 2 moves each, provided we have enough maxChanges. We calculate the total cost and find the minimum among all valid choices of subsegments.

Algorithm

First, identify all indices where nums[i] == 1. Store these indices in a sorted list, let's call it one_indices of size m.
Pre-calculate the prefix sums of the one_indices list to quickly compute the sum of indices in any subsegment.
The core idea is that we will use some number of original ones, say w, and some number of created ones, c, such that w + c = k. The cost for c created ones is 2 * c.
To minimize the cost of gathering w original ones, we must choose a contiguous block of w ones from the one_indices list. Let this block be from index i to j = i + w - 1.
The minimum cost to gather these w ones is to move them all to their median index. This cost can be calculated in O(1) using the pre-calculated prefix sums.
We iterate through all possible contiguous windows [i, j] in one_indices.
For each window, let w = j - i + 1 be the number of original ones used.
The number of changes needed is c = k - w. This is only possible if 0 <= c <= maxChanges.
If the number of changes is valid, calculate the total cost: total_cost = gather_cost(i, w) + 2 * c.
Keep track of the minimum total cost found across all valid windows.
A special case to consider is using only created ones. If k <= maxChanges, a possible answer is 2 * k.

Walkthrough

Let's denote the list of indices of ones as P. The size of this list is m.

First, we can handle the case where we only use 'change' operations. If k <= maxChanges, one possible way to get k ones is to create all of them. This would cost 2 * k moves. This can be an initial value for our minimum cost.

Then, we explore combinations of original ones and created ones. We iterate through all possible start indices i from 0 to m-1 and end indices j from i to m-1. This defines a window P[i...j] of w = j - i + 1 original ones.

For each window, we check if it's a valid choice. We need k-w more ones, which must be created. This requires k-w changes. If k-w > maxChanges or k-w < 0, this window size is not feasible.

If it is feasible, we calculate the cost. The cost has two parts:

cost_changes = 2 * (k - w)
cost_gather: The minimum moves to gather the w original ones from the window P[i...j] to a single point. This point is the median of P[i...j]. The cost can be calculated efficiently using prefix sums of P. Let pref be the prefix sum array. The gathering cost for P[i...j] is (pref[j+1] - pref[median_idx+1]) - (pref[median_idx] - pref[i]) for odd w and (pref[j+1] - pref[median_idx_2]) - (pref[median_idx_1] - pref[i]) for even w, where median indices are calculated based on i and w.

The total cost for the window is cost_gather + cost_changes. We update our overall minimum cost with this value.

1import java.util.ArrayList;2import java.util.List;3 4class Solution {5    public long minimumMoves(int[] nums, int k, int maxChanges) {6        List<Integer> oneIndices = new ArrayList<>();7        for (int i = 0; i < nums.length; i++) {8            if (nums[i] == 1) {9                oneIndices.add(i);10            }11        }12 13        int m = oneIndices.size();14        if (m == 0) {15            return (long)k * 2;16        }17 18        long[] prefixSum = new long[m + 1];19        for (int i = 0; i < m; i++) {20            prefixSum[i + 1] = prefixSum[i] + oneIndices.get(i);21        }22 23        long minMoves = Long.MAX_VALUE;24 25        if (k <= maxChanges) {26            minMoves = (long)k * 2;27        }28 29        for (int w = 1; w <= m; w++) {30            int changesNeeded = k - w;31            if (changesNeeded < 0) break;32            if (changesNeeded > maxChanges) continue;33 34            for (int i = 0; i <= m - w; i++) {35                int j = i + w - 1;36                long gatherCost = getGatherCost(prefixSum, i, j);37                minMoves = Math.min(minMoves, gatherCost + (long)changesNeeded * 2);38            }39        }40 41        return minMoves;42    }43 44    private long getGatherCost(long[] prefixSum, int i, int j) {45        int w = j - i + 1;46        int medianIndexInWindow = (w - 1) / 2;47        int medianAbsIndex = i + medianIndexInWindow;48        49        long leftSum = prefixSum[medianAbsIndex] - prefixSum[i];50        long rightSum = prefixSum[j + 1] - prefixSum[medianAbsIndex + 1];51        long medianVal = (prefixSum[medianAbsIndex + 1] - prefixSum[medianAbsIndex]);52 53        long cost = (medianIndexInWindow * medianVal - leftSum) + (rightSum - (w - 1 - medianIndexInWindow) * medianVal);54        return cost;55    }56}

Complexity

Time

O(m^2), where `m` is the number of ones. The nested loops iterate through all `O(m^2)` possible contiguous subsegments of `one_indices`.

Space

O(m), where `m` is the number of ones in `nums`. This is for storing the indices of ones and their prefix sums.

Trade-offs

Pros

The logic is straightforward and directly models the problem of choosing a contiguous block of original ones.
It's guaranteed to find the optimal solution because it exhaustively checks all valid combinations of contiguous original ones.

Cons

The time complexity is quadratic in the number of ones, which can be too slow if the input array contains many ones.

Solutions

Solution

1class Solution {2public3  long minimumMoves(int[] nums, int k, int maxChanges) {4    int n = nums.length;5    int[] cnt = new int[n + 1];6    long[] s = new long[n + 1];7    for (int i = 1; i <= n; ++i) {8      cnt[i] = cnt[i - 1] + nums[i - 1];9      s[i] = s[i - 1] + i * nums[i - 1];10    }11    long ans = Long.MAX_VALUE;12    for (int i = 1; i <= n; ++i) {13      long t = 0;14      int need = k - nums[i - 1];15      for (int j = i - 1; j <= i + 1; j += 2) {16        if (need > 0 && 1 <= j && j <= n && nums[j - 1] == 1) {17          --need;18          ++t;19        }20      }21      int c = Math.min(need, maxChanges);22      need -= c;23      t += c * 2;24      if (need <= 0) {25        ans = Math.min(ans, t);26        continue;27      }28      int l = 2, r = Math.max(i - 1, n - i);29      while (l <= r) {30        int mid = (l + r) >> 1;31        int l1 = Math.max(1, i - mid), r1 = Math.max(0, i - 2);32        int l2 = Math.min(n + 1, i + 2), r2 = Math.min(n, i + mid);33        int c1 = cnt[r1] - cnt[l1 - 1];34        int c2 = cnt[r2] - cnt[l2 - 1];35        if (c1 + c2 >= need) {36          long t1 = 1L * c1 * i - (s[r1] - s[l1 - 1]);37          long t2 = s[r2] - s[l2 - 1] - 1L * c2 * i;38          ans = Math.min(ans, t + t1 + t2);39          r = mid - 1;40        } else {41          l = mid + 1;42        }43      }44    }45    return ans;46  }47}

1class Solution { public: long long minimumMoves ( vector < int >& nums , int k , int maxChanges ) { int n = nums . size (); vector < int > cnt ( n + 1 , 0 ); vector < long long > s ( n + 1 , 0 ); for ( int i = 1 ; i <= n ; ++ i ) { cnt [ i ] = cnt [ i - 1 ] + nums [ i - 1 ]; s [ i ] = s [ i - 1 ] + 1LL * i * nums [ i - 1 ]; } long long ans = LLONG_MAX ; for ( int i = 1 ; i <= n ; ++ i ) { long long t = 0 ; int need = k - nums [ i - 1 ]; for ( int j = i - 1 ; j <= i + 1 ; j += 2 ) { if ( need > 0 && 1 <= j && j <= n && nums [ j - 1 ] == 1 ) { -- need ; ++ t ; } } int c = min ( need , maxChanges ); need -= c ; t += c * 2 ; if ( need <= 0 ) { ans = min ( ans , t ); continue ; } int l = 2 , r = max ( i - 1 , n - i ); while ( l <= r ) { int mid = ( l + r ) / 2 ; int l1 = max ( 1 , i - mid ), r1 = max ( 0 , i - 2 ); int l2 = min ( n + 1 , i + 2 ), r2 = min ( n , i + mid ); int c1 = cnt [ r1 ] - cnt [ l1 - 1 ]; int c2 = cnt [ r2 ] - cnt [ l2 - 1 ]; if ( c1 + c2 >= need ) { long long t1 = 1LL * c1 * i - ( s [ r1 ] - s [ l1 - 1 ]); long long t2 = s [ r2 ] - s [ l2 - 1 ] - 1LL * c2 * i ; ans = min ( ans , t + t1 + t2 ); r = mid - 1 ; } else { l = mid + 1 ; } } } return ans ; } };

1class Solution:2    def minimumMoves(self, nums: List[int], k: int, maxChanges: int) -> int: n = len(nums) cnt = [0] * (n + 1) s = [0] * (n + 1) for i, x in enumerate(nums, 1): cnt[i] = cnt[i - 1] + x s[i] = s[i - 1] + i * x ans = inf max = lambda x, y: x if x > y else y min = lambda x, y: x if x < y else y for i, x in enumerate(nums, 1): t = 0 need = k - x for j in (i - 1, i + 1): if need > 0 and 1 <= j <= n and nums[j - 1] == 1: need -= 1 t += 1 c = min(need, maxChanges) need -= c t += c * 2 if need <= 0: ans = min(ans, t) continue l, r = 2, max(i - 1, n - i) while l <= r: mid = (l + r) >> 1 l1, r1 = max(1, i - mid), max(0, i - 2) l2, r2 = min(n + 1, i + 2), min(n, i + mid) c1 = cnt[r1] - cnt[l1 - 1] c2 = cnt[r2] - cnt[l2 - 1] if c1 + c2 >= need: t1 = c1 * i - (s[r1] - s[l1 - 1]) t2 = s[r2] - s[l2 - 1] - c2 * i ans = min(ans, t + t1 + t2) r = mid - 1 else: l = mid + 1 return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem