# 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& nums, int k) { const int kIndex = find(nums.begin(), nums.end(), k) - nums.begin(); int ans = 0; unordered_map 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 have a median equal to k. // So, we should add count[0 - balance] and count[1 - balance] to the 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 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 have a median equal to k. // So, we should add count[0 - balance] and count[1 - balance] to the 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 have a median equal to k. # So, we should add count[0 - balance] and count[1 - balance] to the ans. ans += count[-balance] + count[1 - balance] return ans