Skip to content

3086. Minimum Moves to Pick K Ones

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class Solution {
 public:
  long long minimumMoves(vector<int>& nums, int k, int maxChanges) {
    // Dylan has two actions for collecting '1's in a sequence:
    // Action 1: Put a '1' next to him and pick it up.
    //           The cost is 2.
    // Action 2: Swap a '1' towards him and collect it.
    //           The cost equals the distance to the '1'.
    //
    // To minimize the swapping cost, Dylan can use a sliding window strategy,
    // selecting the optimal position (middle '1' in the window) for efficient
    // collection. The window's size is crucial:

    // The minimum window size: min(0, k - maxChanges), ensuring the window
    // isn't too small.
    // The maximum window size: min(k, minOnesByTwo + 3, the number of ones),
    // preventing overly ambitious swaps.
    //
    // Note that if needing to move a '1' beyond `minOnesByTwo + 3`, it's
    // cheaper to use Action 1.

    // At most three indices, (dylanIndex - 1, dylanIndex, dylanIndex + 1), have
    // a distance <= 1 from dylanIndex, implying that we'll be taking at most
    // `maxOnesByTwo + 3` using Action 2. Any more Action 2 is not optimal and
    // should be replaced with Action 1.
    constexpr int kNumOfIndicesWithinOneDistance = 3;
    long ans = LONG_MAX;
    vector<int> oneIndices;  // the indices of 1s
    vector<long> prefix{0};  // the accumulated indices of 1s

    for (int i = 0; i < nums.size(); ++i)
      if (nums[i] == 1)
        oneIndices.push_back(i);

    for (const int oneIndex : oneIndices)
      prefix.push_back(prefix.back() + oneIndex);

    const int minOnesByTwo = max(0, k - maxChanges);
    const int maxOnesByTwo =
        min({k, minOnesByTwo + kNumOfIndicesWithinOneDistance,
             static_cast<int>(oneIndices.size())});

    for (int onesByTwo = minOnesByTwo; onesByTwo <= maxOnesByTwo; ++onesByTwo)
      for (int l = 0; l + onesByTwo < prefix.size(); ++l) {
        const int r = l + onesByTwo;  // Collect 1s in oneIndices[l - 1..r - 1].
        const long cost1 = (k - onesByTwo) * 2;
        const long cost2 = (prefix[r] - prefix[(l + r) / 2]) -
                           (prefix[(l + r + 1) / 2] - prefix[l]);
        ans = min(ans, cost1 + cost2);
      }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
class Solution {
  public long minimumMoves(int[] nums, int k, int maxChanges) {
    // Dylan has two actions for collecting '1's in a sequence:
    // Action 1: Put a '1' next to him and pick it up.
    //           The cost is 2.
    // Action 2: Swap a '1' towards him and collect it.
    //           The cost equals the distance to the '1'.
    //
    // To minimize the swapping cost, Dylan can use a sliding window strategy,
    // selecting the optimal position (middle '1' in the window) for efficient
    // collection. The window's size is crucial:

    // The minimum window size: min(0, k - maxChanges), ensuring the window
    // isn't too small.
    // The maximum window size: min(k, minOnesByTwo + 3, the number of ones),
    // preventing overly ambitious swaps.
    //
    // Note that if needing to move a '1' beyond `minOnesByTwo + 3`, it's
    // cheaper to use Action 1.

    // At most three indices, (dylanIndex - 1, dylanIndex, dylanIndex + 1), have
    // a distance <= 1 from dylanIndex, implying that we'll be taking at most
    // `maxOnesByTwo + 3` using Action 2. Any more Action 2 is not optimal and
    // should be replaced with Action 1.
    final int NUM_OF_INDICES_WITHIN_ONE_DISTANCE = 3;
    long ans = Long.MAX_VALUE;
    List<Integer> oneIndices = new ArrayList<>(); // the indices of 1s
    List<Long> prefix = new ArrayList<>();        // the accumulated indices of 1s
    prefix.add(0L);

    for (int i = 0; i < nums.length; ++i)
      if (nums[i] == 1)
        oneIndices.add(i);

    for (final int oneIndex : oneIndices)
      prefix.add(prefix.get(prefix.size() - 1) + oneIndex);

    final int minOnesByTwo = Math.max(0, k - maxChanges);
    final int maxOnesByTwo =
        Math.min(k, Math.min(minOnesByTwo + NUM_OF_INDICES_WITHIN_ONE_DISTANCE, oneIndices.size()));

    for (int onesByTwo = minOnesByTwo; onesByTwo <= maxOnesByTwo; ++onesByTwo)
      for (int l = 0; l + onesByTwo < prefix.size(); ++l) {
        final int r = l + onesByTwo; // Collect 1s in oneIndices[l - 1..r - 1].
        final long cost1 = (k - onesByTwo) * 2;
        final long cost2 = (prefix.get(r) - prefix.get((l + r) / 2)) -
                           (prefix.get((l + r + 1) / 2) - prefix.get(l));
        ans = Math.min(ans, cost1 + cost2);
      }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Solution:
  def minimumMoves(self, nums: list[int], k: int, maxChanges: int) -> int:
    # Dylan has two actions for collecting '1's in a sequence:
    # Action 1: Put a '1' next to him and pick it up.
    #           The cost is 2.
    # Action 2: Swap a '1' towards him and collect it.
    #           The cost equals the distance to the '1'.
    #
    # To minimize the swapping cost, Dylan can use a sliding window strategy,
    # selecting the optimal position (middle '1' in the window) for efficient
    # collection. The window's size is crucial:

    # The minimum window size: min(0, k - maxChanges), ensuring the window
    # isn't too small.
    # The maximum window size: min(k, minOnesByTwo + 3, the number of ones),
    # preventing overly ambitious swaps.
    #
    # Note that if needing to move a '1' beyond `minOnesByTwo + 3`, it's
    # cheaper to use Action 1.

    # At most three indices, (dylanIndex - 1, dylanIndex, dylanIndex + 1), have
    # a distance <= 1 from dylanIndex, implying that we'll be taking at most
    # `maxOnesByTwo + 3` using Action 2. Any more Action 2 is not optimal and
    # should be replaced with Action 1.
    NUM_OF_INDICES_WITHIN_ONE_DISTANCE = 3
    ans = math.inf
    oneIndices = [i for i, num in enumerate(nums) if num == 1]
    prefix = list(itertools.accumulate(oneIndices, initial=0))

    minOnesByTwo = max(0, k - maxChanges)
    maxOnesByTwo = min(
        k, minOnesByTwo + NUM_OF_INDICES_WITHIN_ONE_DISTANCE, len(oneIndices))

    for onesByTwo in range(minOnesByTwo, maxOnesByTwo + 1):
      for l in range(len(prefix) - onesByTwo):
        r = l + onesByTwo  # Collect 1s in oneIndices[l - 1..r - 1].
        cost1 = (k - onesByTwo) * 2
        cost2 = ((prefix[r] - prefix[(l + r) // 2]) -
                 (prefix[(l + r + 1) // 2] - prefix[l]))
        ans = min(ans, cost1 + cost2)

    return ans