Skip to content

2607. Make K-Subarray Sums Equal 👍

  • Time: C++: $O(n)$, Java/Python: $O(n\log 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
class Solution {
 public:
  long long makeSubKSumEqual(vector<int>& arr, int k) {
    // If the sum of each subarray of length k is equal, then `arr` must have a
    // repeated pattern of size k. e.g. arr = [1, 2, 3, ...] and k = 3, to have
    // sum([1, 2, 3)] == sum([2, 3, x]), x must be 1. Therefore, arr[i] ==
    // arr[(i + k) % n] for every i.
    const int n = arr.size();
    long ans = 0;
    vector<bool> seen(n);

    for (int i = 0; i < n; ++i) {
      vector<int> groups;
      int j = i;
      while (!seen[j]) {
        groups.push_back(arr[j]);
        seen[j] = true;
        j = (j + k) % n;
      }
      nth_element(groups.begin(), groups.begin() + groups.size() / 2,
                  groups.end());
      for (const int num : groups)
        ans += abs(num - groups[groups.size() / 2]);
    }

    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
class Solution {
  public long makeSubKSumEqual(int[] arr, int k) {
    // If the sum of each subarray of length k is equal, then `arr` must have a
    // repeated pattern of size k. e.g. arr = [1, 2, 3, ...] and k = 3, to have
    // sum([1, 2, 3)] == sum([2, 3, x]), x must be 1. Therefore, arr[i] ==
    // arr[(i + k) % n] for every i.
    final int n = arr.length;
    long ans = 0;
    boolean[] seen = new boolean[n];

    for (int i = 0; i < n; ++i) {
      List<Integer> groups = new ArrayList<>();
      int j = i;
      while (!seen[j]) {
        groups.add(arr[j]);
        seen[j] = true;
        j = (j + k) % n;
      }
      Collections.sort(groups);
      for (final int num : groups)
        ans += Math.abs(num - groups.get(groups.size() / 2));
    }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
  def makeSubKSumEqual(self, arr: list[int], k: int) -> int:
    # If the sum of each subarray of length k is equal, then `arr` must have a
    # repeated pattern of size k. e.g. arr = [1, 2, 3, ...] and k = 3, to have
    # sum([1, 2, 3)] == sum([2, 3, x]), x must be 1. Therefore, arr[i] ==
    # arr[(i + k) % n] for every i.
    n = len(arr)
    ans = 0
    seen = [0] * n

    for i in range(n):
      groups = []
      j = i
      while not seen[j]:
        groups.append(arr[j])
        seen[j] = True
        j = (j + k) % n
      groups.sort()
      for num in groups:
        ans += abs(num - groups[len(groups) // 2])

    return ans