Create count-subarrays-with-median-k.cpp

Author: kamyu104

C++/count-subarrays-with-median-k.cpp

@@ -0,0 +1,21 @@
+// Time:  O(n)
+// Space: O(n)
+
+// freq table, prefix sum
+class Solution {
+public:
+    int countSubarrays(vector<int>& nums, int k) {
+        const int idx = distance(cbegin(nums), find(cbegin(nums), cend(nums), k));
+        unordered_map<int, int> lookup;
+        for (int i = idx, curr = 0; i >= 0; --i) {
+            curr += nums[i] == k ? 0 : (nums[i] < k ? -1 : +1);
+            ++lookup[curr];
+        }
+        int result = 0;
+        for (int i = idx, curr = 0; i < size(nums); ++i) {
+            curr += nums[i] == k ? 0 : (nums[i] < k ? -1 : +1);
+            result += lookup[-curr] + lookup[-(curr - 1)];
+        }
+        return result;
+    }
+};