diff --git a/19_HashMap/16. Count Subarrays With Median K.py b/19_HashMap/16. Count Subarrays With Median K.py new file mode 100644 index 00000000..5f5e58e5 --- /dev/null +++ b/19_HashMap/16. Count Subarrays With Median K.py @@ -0,0 +1,20 @@ +# https://leetcode.com/problems/count-subarrays-with-median-k/ +# https://youtu.be/QZzDioqkRhU + +class Solution: + def countSubarrays(self, nums: List[int], k: int) -> int: + pos = nums.index(k) + cnt = defaultdict(int) + cnt[0] = 1 + c = 0 + for i in range(pos + 1, len(nums)): + c += 1 if nums[i] > k else -1 + cnt[c] += 1 + + ans = cnt[0] + cnt[1] + c = 0 + for i in range(pos - 1, -1, -1): + c += 1 if nums[i] < k else -1 + ans += cnt[c] + cnt[c + 1] + # print(cnt) + return ans