Skip to content

2488. Count Subarrays With Median K 👍

  • 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
class Solution {
 public:
  int countSubarrays(vector<int>& nums, int k) {
    const int kIndex = find(nums.begin(), nums.end(), k) - nums.begin();
    int ans = 0;
    unordered_map<int, int> count;

    for (int i = kIndex, balance = 0; i >= 0; --i) {
      if (nums[i] < k)
        --balance;
      else if (nums[i] > k)
        ++balance;
      ++count[balance];
    }

    for (int i = kIndex, balance = 0; i < nums.size(); ++i) {
      if (nums[i] < k)
        --balance;
      else if (nums[i] > k)
        ++balance;
      // The subarray that has balance == 0 or 1 having median equal to k.
      // So, add count[0 - balance] and count[1 - balance] to `ans`.
      ans += count[-balance] + count[1 - balance];
    }

    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
class Solution {
  public int countSubarrays(int[] nums, int k) {
    final int kIndex = find(nums, k);
    int ans = 0;
    Map<Integer, Integer> count = new HashMap<>();

    for (int i = kIndex, balance = 0; i >= 0; --i) {
      if (nums[i] < k)
        --balance;
      else if (nums[i] > k)
        ++balance;
      count.merge(balance, 1, Integer::sum);
    }

    for (int i = kIndex, balance = 0; i < nums.length; ++i) {
      if (nums[i] < k)
        --balance;
      else if (nums[i] > k)
        ++balance;
      // The subarray that has balance == 0 or 1 having median equal to k.
      // So, add count[0 - balance] and count[1 - balance] to `ans`.
      ans += count.getOrDefault(-balance, 0) + count.getOrDefault(1 - balance, 0);
    }

    return ans;
  }

  private int find(int[] nums, int k) {
    for (int i = 0; i < nums.length; ++i)
      if (nums[i] == k)
        return i;
    throw new IllegalArgumentException();
  }
}

 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
class Solution:
  def countSubarrays(self, nums: list[int], k: int) -> int:
    kIndex = nums.index(k)
    ans = 0
    count = collections.Counter()

    balance = 0
    for i in range(kIndex, -1, -1):
      if nums[i] < k:
        balance -= 1
      elif nums[i] > k:
        balance += 1
      count[balance] += 1

    balance = 0
    for i in range(kIndex, len(nums)):
      if nums[i] < k:
        balance -= 1
      elif nums[i] > k:
        balance += 1
      # The subarray that has balance == 0 or 1 having median equal to k.
      # So, add count[0 - balance] and count[1 - balance] to `ans`.
      ans += count[-balance] + count[1 - balance]

    return ans