Skip to content

1703. Minimum Adjacent Swaps for K Consecutive 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
class Solution {
 public:
  int minMoves(vector<int>& nums, int k) {
    vector<int> ones;

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

    // Returns the median index of [i..i + k).
    auto getMedIndex = [&](int i) { return (i + (i + k - 1)) / 2; };

    // Calculate the first group: window[0] = A[0..k).
    const int median = ones[getMedIndex(0)];
    int moves = 0;
    for (int i = 0; i < k; ++i)
      moves += abs(ones[i] - median);

    int ans = moves;

    for (int i = 1; i <= ones.size() - k; ++i) {
      const int oldMedianIndex = ones[getMedIndex(i - 1)];
      const int newMedianIndex = ones[getMedIndex(i)];
      if (k % 2 == 1)
        moves += newMedianIndex - oldMedianIndex;
      moves -= newMedianIndex - ones[i - 1];
      moves += ones[i + k - 1] - newMedianIndex;
      ans = min(ans, moves);
    }

    auto nThSum = [&](int n) { return n * (n + 1) / 2; };
    return ans - nThSum((k - 1) / 2) - nThSum(k / 2);
  }
};

 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
class Solution {
  public int minMoves(int[] nums, int k) {
    List<Integer> ones = new ArrayList<>();

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

    final int median = ones.get(getMedIndex(0, k));
    int moves = 0;
    for (int i = 0; i < k; ++i)
      moves += Math.abs(ones.get(i) - median);

    int ans = moves;

    for (int i = 1; i <= ones.size() - k; ++i) {
      final int oldMedianIndex = ones.get(getMedIndex(i - 1, k));
      final int newMedianIndex = ones.get(getMedIndex(i, k));
      if (k % 2 == 1)
        moves += newMedianIndex - oldMedianIndex;
      moves -= newMedianIndex - ones.get(i - 1);
      moves += ones.get(i + k - 1) - newMedianIndex;
      ans = Math.min(ans, moves);
    }

    return ans - nThSum((k - 1) / 2) - nThSum(k / 2);
  }

  // Returns the median index of [i..i + k).
  private int getMedIndex(int i, int k) {
    return (i + (i + k - 1)) / 2;
  }

  // Returns 1 + 2 + ... + n
  private int nThSum(int n) {
    return n * (n + 1) / 2;
  }
}