3086. Minimum Moves to Pick K Ones

Problem Description

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 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.

Intuition

The problem requires finding an optimal strategy for Alice to pick up k ones efficiently. The challenge lies in the movement and change constraints.

Standing Position Advantage: Start by considering the possibility of picking up a 1 at Alice's initial standing position without making a move. This serves as an immediate advantage.
Adjacent Movement: Next, check the positions immediately adjacent to Alice. If she can use existing ones in these positions to her advantage, it minimizes changes required and saves moves. This can be accomplished by swapping if adjacent positions have a 1.
Making Zero into One: The problem allows for converting up to maxChanges zeros to ones, providing further opportunities to pick up ones. Using these changes strategically can significantly aid the process.
Binary Search for Optimal Interval: For any remaining ones needed, consider an expanded interval around Alice's position. This involves searching for an optimal interval that contains sufficient ones, minimally disrupting the structure of the array.

By combining greedy strategies (using the ones immediately available) with calculated adjustments (changing zeros to ones and searching for clusters of ones), the problem can be broken down into manageable parts that lead to an efficient solution.

Learn more about Greedy, Prefix Sum and Sliding Window patterns.

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Solution Approach

The solution uses a combination of greedy methods, prefix sums, and binary search to achieve the goal efficiently.

Prefix Sum Setup:

First, calculate prefix sums to quickly compute the number of ones and their indices within any subarray. This helps in making fast range queries for how many 1s exist in slices of the array.

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

Greedy Enumeration:
- Iterate over each possible starting position aliceIndex. At each position, devise a strategy to gather exactly k ones.
Local Optimization:
- First, check if aliceIndex has a 1 to pick up without a move. Then consider swapping adjacent 0s for 1s when possible. This local adjustment minimizes moves and leverages nearby resources.
Changing Zero to One:
- Utilize up to maxChanges times to convert zeros into ones if it effectively reduces the distance to gather necessary ones. Plan these changes carefully since each counts as a potential future move.

Binary Search for Optimal Range:

Use binary search to determine if an expanded range of indices can be leveraged to gather remaining ones efficiently. The binary search helps find the smallest interval around i that contains enough ones to meet the k requirement.

l, r = 2, max(i - 1, n - i)  # Initial range setup
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]
    # Check if the number of ones available within the range is sufficient
    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

By strategically combining these patterns—leveraging prefix sums for fast queries, using greedy local optimizations, and deploying binary search to determine effective ranges—the solution finds the minimum moves required to achieve Alice's target efficiently.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's consider a small example to illustrate the solution approach. Suppose nums = [1, 0, 0, 1, 0, 1], k = 2, and maxChanges = 1.

Step 1: Initial Setup

Calculate Prefix Sums:
- We initialize two arrays, cnt and s, of length n+1 (where n is the length of nums).
- cnt[i]: The count of 1s up to the i-th index.
- s[i]: The sum of indices where there is a 1.
Result:
```
cnt = [0, 1, 1, 1, 2, 2, 3]
s = [0, 1, 1, 1, 5, 5, 11]
```

Step 2: Greedy Enumeration and Local Optimization

Iterate Over Each Possible Starting Position:

Example with aliceIndex = 0:
- nums[0] == 1, so Alice picks it up without a move. We now need k-1 = 1 more 1.
- Adjacent Movement:
  - Check adjacent indices for swapping. aliceIndex is 0, so the next index is 1. nums[1] == 0, no immediate swap available.
- Use maxChanges:
  - We can change nums[1] from 0 to 1 using one of our allowed changes. Now we have 1s at indices 1 and 3.

Step 3: Binary Search for Optimal Interval

Binary Search for Remaining 1s:

With maxChanges, check if enough 1s are collected within new range:
- Binary search possible intervals. We are seeking to make optimal use of the allowed change.
- Attempt to cover enough array portion with ones.
Objective Check:

Using maxChanges, Alice successfully gathers two 1s. The first one was picked immediately, and using change, we secure the next needed one effectively.

Summary:

Alice starts by picking existing 1s and uses her allowed changes efficiently.
The algorithm reduces the region to examine via binary search for intervals meeting the requirement.
Ultimately, Alice accomplishes her goal using minimal steps and strategic placement of changes.

Solution Implementation

1from typing import List
2from math import inf
3
4class Solution:
5    def minimumMoves(self, nums: List[int], k: int, maxChanges: int) -> int:
6        n = len(nums)
7      
8        # Cumulative arrays to keep track of the count of 1s and weighted sum.
9        count_ones = [0] * (n + 1)
10        weighted_sum = [0] * (n + 1)
11      
12        # Calculate cumulative count of 1s and their weighted sum.
13        for i, x in enumerate(nums, 1):
14            count_ones[i] = count_ones[i - 1] + x
15            weighted_sum[i] = weighted_sum[i - 1] + i * x
16      
17        min_moves = inf  # Initialize the minimum moves as infinity 
18      
19        # Iterate over each element in the list.
20        for i, x in enumerate(nums, 1):
21            total_moves = 0
22            needed_ones = k - x
23          
24            # Check the adjacent positions for existing 1s.
25            for j in (i - 1, i + 1):
26                if needed_ones > 0 and 1 <= j <= n and nums[j - 1] == 1:
27                    needed_ones -= 1
28                    total_moves += 1
29          
30            # Cap the number of changes using maxChanges
31            changes_made = min(needed_ones, maxChanges)
32            needed_ones -= changes_made
33            total_moves += changes_made * 2
34          
35            if needed_ones <= 0:
36                min_moves = min(min_moves, total_moves)
37                continue
38          
39            # Binary search to find the optimal position for making k consecutive 1s
40            left, right = 2, max(i - 1, n - i)
41            while left <= right:
42                mid = (left + right) // 2
43                l1, r1 = max(1, i - mid), max(0, i - 2)
44                l2, r2 = min(n + 1, i + 2), min(n, i + mid)
45              
46                # Count the ones in the defined left and right segments
47                count_left = count_ones[r1] - count_ones[l1 - 1]
48                count_right = count_ones[r2] - count_ones[l2 - 1]
49              
50                # If sufficient 1s are found in these segments
51                if count_left + count_right >= needed_ones:
52                    # Calculate the minimum moves needed by calculating the potential positions
53                    moves_left = count_left * i - (weighted_sum[r1] - weighted_sum[l1 - 1])
54                    moves_right = (weighted_sum[r2] - weighted_sum[l2 - 1]) - count_right * i
55                    min_moves = min(min_moves, total_moves + moves_left + moves_right)
56                    right = mid - 1
57                else:
58                    left = mid + 1
59      
60        return min_moves
61

1class Solution {
2    public long minimumMoves(int[] nums, int k, int maxChanges) {
3        int n = nums.length;
4      
5        // Arrays to store cumulative sums
6        int[] countPrefix = new int[n + 1]; // Prefix sum of 1s
7        long[] weightedSumPrefix = new long[n + 1]; // Weighted prefix sum
8      
9        // Calculate prefix sums
10        for (int i = 1; i <= n; ++i) {
11            countPrefix[i] = countPrefix[i - 1] + nums[i - 1];
12            weightedSumPrefix[i] = weightedSumPrefix[i - 1] + i * nums[i - 1];
13        }
14      
15        long minimumMoves = Long.MAX_VALUE;
16      
17        // Iterate over each position
18        for (int i = 1; i <= n; ++i) {
19            long currentMoves = 0;
20            int neededOnes = k - nums[i - 1];
21          
22            // Check immediate neighbors for available ones
23            for (int j = i - 1; j <= i + 1; j += 2) {
24                if (neededOnes > 0 && 1 <= j && j <= n && nums[j - 1] == 1) {
25                    --neededOnes;
26                    ++currentMoves;
27                }
28            }
29          
30            // Use the maxChanges to minimize needed ones
31            int changesUsed = Math.min(neededOnes, maxChanges);
32            neededOnes -= changesUsed;
33            currentMoves += changesUsed * 2;
34          
35            if (neededOnes <= 0) {
36                // If no more ones are needed, update the minimum moves
37                minimumMoves = Math.min(minimumMoves, currentMoves);
38                continue;
39            }
40          
41            // Binary search to find the minimum range to collect needed ones
42            int left = 2, right = Math.max(i - 1, n - i);
43            while (left <= right) {
44                int mid = (left + right) >> 1;
45                int left1 = Math.max(1, i - mid), right1 = Math.max(0, i - 2);
46                int left2 = Math.min(n + 1, i + 2), right2 = Math.min(n, i + mid);
47              
48                int count1 = countPrefix[right1] - countPrefix[left1 - 1];
49                int count2 = countPrefix[right2] - countPrefix[left2 - 1];
50              
51                if (count1 + count2 >= neededOnes) {
52                    long moves1 = 1L * count1 * i - (weightedSumPrefix[right1] - weightedSumPrefix[left1 - 1]);
53                    long moves2 = weightedSumPrefix[right2] - weightedSumPrefix[left2 - 1] - 1L * count2 * i;
54                    minimumMoves = Math.min(minimumMoves, currentMoves + moves1 + moves2);
55                    right = mid - 1; // Narrow down the range
56                } else {
57                    left = mid + 1;
58                }
59            }
60        }
61        return minimumMoves;
62    }
63}
64

1class Solution {
2public:
3    long long minimumMoves(vector<int>& nums, int k, int maxChanges) {
4        int n = nums.size();
5        vector<int> cnt(n + 1, 0); // Count of numbers seen up to each position
6        vector<long long> sum(n + 1, 0); // Prefix sum of weighted indices
7
8        // Calculate the prefix count and prefix weighted sum
9        for (int i = 1; i <= n; ++i) {
10            cnt[i] = cnt[i - 1] + nums[i - 1];
11            sum[i] = sum[i - 1] + 1LL * i * nums[i - 1];
12        }
13
14        long long answer = LLONG_MAX; // Initialize the answer to maximum possible value
15
16        // Iterate over each element in nums
17        for (int i = 1; i <= n; ++i) {
18            long long movesRequired = 0;
19            int need = k - nums[i - 1];
20
21            // Check left and right elements within one position
22            for (int j = i - 1; j <= i + 1; j += 2) {
23                if (need > 0 && 1 <= j && j <= n && nums[j - 1] == 1) {
24                    --need;
25                    ++movesRequired;
26                }
27            }
28
29            // Change as many zeros to ones as allowed by maxChanges
30            int changes = min(need, maxChanges);
31            need -= changes;
32            movesRequired += changes * 2; // Changing zeros costs double moves
33
34            // If remaining need is zero, update answer and continue
35            if (need <= 0) {
36                answer = min(answer, movesRequired);
37                continue;
38            }
39
40            // Use binary search to find minimum moves required in larger range
41            int left = 2, right = max(i - 1, n - i);
42
43            while (left <= right) {
44                int middle = (left + right) / 2;
45                // Calculate range ends for left and right neighbors
46                int left1 = max(1, i - middle), right1 = max(0, i - 2);
47                int left2 = min(n + 1, i + 2), right2 = min(n, i + middle);
48
49                // Count how many 1s are in both ranges
50                int count1 = cnt[right1] - cnt[left1 - 1];
51                int count2 = cnt[right2] - cnt[left2 - 1];
52
53                // Check if the total 1s can satisfy the need
54                if (count1 + count2 >= need) {
55                    // Calculate the cost to move them
56                    long long t1 = 1LL * count1 * i - (sum[right1] - sum[left1 - 1]);
57                    long long t2 = sum[right2] - sum[left2 - 1] - 1LL * count2 * i;
58                    answer = min(answer, movesRequired + t1 + t2); // Update answer if minimum
59                    right = middle - 1; // Explore smaller range
60                } else {
61                    left = middle + 1; // Explore larger range
62                }
63            }
64        }
65
66        return answer;
67    }
68};
69

1function minimumMoves(nums: number[], k: number, maxChanges: number): number {
2    const n = nums.length;
3    const cnt = Array(n + 1).fill(0); // Cumulative array for number of 1's
4    const s = Array(n + 1).fill(0); // Cumulative array for index-times-nums[i]
5
6    // Populate cumulative arrays
7    for (let 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
12    let minimumMovesResult = Infinity; // Start with the highest possible value
13
14    // Evaluate each position as a potential starting index for the segment
15    for (let i = 1; i <= n; i++) {
16        let moveCount = 0;
17        let neededOnes = k - nums[i - 1];
18
19        // Check neighboring positions
20        for (const neighborIndex of [i - 1, i + 1]) {
21            if (neededOnes > 0 && 1 <= neighborIndex && neighborIndex <= n && nums[neighborIndex - 1] === 1) {
22                neededOnes--;
23                moveCount++;
24            }
25        }
26
27        // Adjust needed ones by using maxChanges
28        const changesUsed = Math.min(neededOnes, maxChanges);
29        neededOnes -= changesUsed;
30        moveCount += changesUsed * 2;
31
32        if (neededOnes <= 0) {
33            minimumMovesResult = Math.min(minimumMovesResult, moveCount);
34            continue;
35        }
36
37        let left = 2,
38            right = Math.max(i - 1, n - i);
39
40        // Binary search to find optimal segment length
41        while (left <= right) {
42            const mid = (left + right) >> 1;
43            const [leftStart, leftEnd] = [Math.max(1, i - mid), Math.max(0, i - 2)];
44            const [rightStart, rightEnd] = [Math.min(n + 1, i + 2), Math.min(n, i + mid)];
45
46            const leftOnes = cnt[leftEnd] - cnt[leftStart - 1];
47            const rightOnes = cnt[rightEnd] - cnt[rightStart - 1];
48
49            if (leftOnes + rightOnes >= neededOnes) {
50                const leftCost = leftOnes * i - (s[leftEnd] - s[leftStart - 1]);
51                const rightCost = s[rightEnd] - s[rightStart - 1] - rightOnes * i;
52                minimumMovesResult = Math.min(minimumMovesResult, moveCount + leftCost + rightCost);
53                right = mid - 1;
54            } else {
55                left = mid + 1;
56            }
57        }
58    }
59
60    return minimumMovesResult;
61}
62

Time and Space Complexity

The time complexity of the code is O(n × log n). This is due to the binary search step that is performed for each element i in the nums list, where the length of nums is n. Specifically, for each element, a binary search is used to find a minimal range which results in the logarithmic factor in combination with the linear pass over the list.

The space complexity is O(n). This is because two lists cnt and s are used, each with a size of n+1, which are used to store cumulative sums and their indices. These lists are used to quickly query subarray sums and are directly proportional to the size of the input array nums.

Learn more about how to find time and space complexity quickly.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Given a sorted array of integers and an integer called target, find the element that equals to the target and return its index. Select the correct code that fills the ___ in the given code snippet.

1def binary_search(arr, target):
2    left, right = 0, len(arr) - 1
3    while left ___ right:
4        mid = (left + right) // 2
5        if arr[mid] == target:
6            return mid
7        if arr[mid] < target:
8            ___ = mid + 1
9        else:
10            ___ = mid - 1
11    return -1
12

1public static int binarySearch(int[] arr, int target) {
2    int left = 0;
3    int right = arr.length - 1;
4
5    while (left ___ right) {
6        int mid = left + (right - left) / 2;
7        if (arr[mid] == target) return mid;
8        if (arr[mid] < target) {
9            ___ = mid + 1;
10        } else {
11            ___ = mid - 1;
12        }
13    }
14    return -1;
15}
16

1function binarySearch(arr, target) {
2    let left = 0;
3    let right = arr.length - 1;
4
5    while (left ___ right) {
6        let mid = left + Math.trunc((right - left) / 2);
7        if (arr[mid] == target) return mid;
8        if (arr[mid] < target) {
9            ___ = mid + 1;
10        } else {
11            ___ = mid - 1;
12        }
13    }
14    return -1;
15}
16

<=, left, right

<, left, right

<=, right, left

<, right, left

I don't know