2488. Count Subarrays With Median K
Problem Description
You have been given an array named nums
with a size n
, which includes distinct integers ranging from 1
to n
. You are also given a positive integer k
. Your task is to count how many non-empty subarrays of nums
have a median that is equal to k
.
Remember, the median of an array is defined as the middle element once the array is sorted in ascending order. For an array with an even number of elements, there isn’t a single middle element, so in that case, you take the left middle element as the median.
A subarray is any sequence of consecutive elements from the array. It can be as short as one element, or as long as the entire array.
To illustrate, consider the array [2, 3, 1, 4]
. After sorting, this array becomes [1, 2, 3, 4]
. The median is the left middle element, which is 2
for this array. For the array [8, 4, 3, 5, 1]
, after sorting it becomes [1, 3, 4, 5, 8]
. The median for this array is 4
since there is an odd number of elements.
Your objective is to find and return the number of these subarrays in which the median is exactly k
.
Intuition
Let's try to decode the provided solution to this problem. The main idea here is to use the index of k
in nums
as an anchor point. Since the problem tells us nums
consists of distinct integers from 1
to n
, we know k
will appear exactly once in nums
.
We can search for the contiguous subarrays where k
is the median. If k
is the median of a subarray, for the subarray to remain valid when adding elements to the left or right, we must maintain a balance or near-balance in the number of elements that are less than k
and the number of elements greater than k
.
While iterating through the elements, we keep track of this balance with a counter x
, which we increment if we add an element greater than k
to our subarray, and decrement if the element is less than k
. The counter value represents the difference between the count of elements greater than k
and less than k
to the right of k
.
The solution approach calculates this balance for the elements to the right and left of k
separately, then combines these results. For the elements on the right of k
, we just need to keep track of instances where this balance is 0 or 1. For the elements on the left of k
, we also have to consider the complementary counts from the right side to ensure that k
remains the median when combining subarrays from both sides.
Essentially, the algorithm loops twice: once over the elements to the right of k
's index and once over the elements to the left. In each loop, it maintains a count (cnt
) of occurrences of each balance value. This count is key to finding the number of valid subarrays when we combine elements from both sides. The final answer is the sum of all valid subarray counts.
Learn more about Prefix Sum patterns.
Solution Approach
The implementation of the solution revolves around a two-part process. In the first part, we find the index i
of the integer k
within the array nums
using nums.index(k)
. This is crucial as we will compute the balance of larger and smaller elements in the subarrays with respect to k
.
To assist in counting, a Counter
object named cnt
is instantiated. This data structure is a dictionary subclass from Python's collections
module that counts hashable objects.
The core part of the algorithm operates on the premise that when we add an element greater than k
, it could potentially disrupt the balance required for k
to be the median. Thus, every time we encounter such an element, we increase our balance counter x
. Conversely, an element less than k
contributes toward maintaining the median, and thus we decrease the counter x
by one.
Analyzing the Right of k
We create a subarray of nums
starting from the element just after k
with nums[i + 1 :]
. With a for-loop over this subarray, for every element v
that is greater than k
, we increment x
and for every element less than k
, we decrement x
. The conditions 0 <= x <= 1
ensure that the median of the new subarray including v
remains k
(as k
should not be overtaken by a higher number of elements that are either greater or smaller). For each balance situation satisfying this condition, we increase the solution count ans
. The cnt
counter keeps track of how many times each x
value occurred.
Analyzing the Left of k
Once the right side is processed, we reset x
to 0
and run a second for-loop, iterating over the elements to the left of k
. The loop decrements the index j
each time, moving leftwards. For each element nums[j]
, x
is adjusted in the same way (increment for elements greater than k
, decrement for elements less than k
).
For each balance value x
, two additions are made to the ans
:
0 <= x <= 1
is again checked, and if true,ans
is incremented.- The corresponding counts for
-x
and-x + 1
in thecnt
Counter (the balance counts from the right side) are added toans
.
These additions account for the subarrays extending both left and right from element k
, whose medians are still k
.
Finally, the ans
which holds the count of all qualifying subarrays is returned.
This solution leverages the balance theory, counter collections, indexing, and iteration in reverse to efficiently find all the subarrays with k
as their median.
Ready to land your dream job?
Unlock your dream job with a 2-minute evaluator for a personalized learning plan!
Start EvaluatorExample Walkthrough
Let's walk through a small example to illustrate the solution approach. Consider the array nums = [3, 1, 4, 2]
with n = 4
and let's find the number of subarrays where the median equals k = 2
.
First, we need to find the index of k
in nums
; k
is at index i = 3
.
Now, let's analyze the elements to the right of k
, which in this case, are none because k
is at the last index. So, we move on to the next part of the process.
Next, we will analyze the elements to the left of k
. We iterate from index i - 1
to 0
(that is, from right to left). We initialize our balance counter x
to 0
and a Counter to keep track of the occurrences of x
.
- At index
i - 1
(index 2), the value is4
(greater thank
). We incrementx
to 1 (x = x + 1
). There are no elements to the right, so we don't update our answer yet. - At index
i - 2
(index 1), the value is1
(less thank
). We decrementx
to 0 (x = x - 1
). If we were to choose the subarray[3, 1, 4, 2]
up to now, the sorted version would be[1, 2, 3, 4]
, andk
would be the median. We increment our answerans
by 1. - At index
i - 3
(index 0), the value is3
(greater thank
). We incrementx
to 1 again. The subarray[1, 4, 2]
hask
as the median after sorting to[1, 2, 4]
.ans
is incremented by 1.
Finishing this process, we have found 2
subarrays where k
is the median: [3, 1, 4, 2]
and [1, 4, 2]
. Therefore, by following the steps outlined in the solution approach, we would return 2
as the final answer.
Solution Implementation
1from collections import Counter
2
3class Solution:
4 def count_subarrays(self, nums: List[int], k: int) -> int:
5 # Find the index of the first occurrence of k in nums.
6 index_of_k = nums.index(k)
7
8 # Initialize a counter to count occurrences of certain balances of values vs k.
9 balance_counter = Counter()
10
11 # Start with one subarray that includes k itself.
12 total_subarrays = 1
13
14 # 'balance' keeps track of the difference between the count of numbers greater than k
15 # and count of numbers less than k.
16 balance = 0
17
18 # Look at the elements to the right of k in nums.
19 for value in nums[index_of_k + 1:]:
20 # If the number is greater than k, increment the balance; if less, decrement.
21 balance += 1 if value > k else -1
22
23 # Increment the count for subarrays where balance is 0 (equal numbers of both sides) or 1.
24 total_subarrays += 0 <= balance <= 1
25
26 # Update the counter for current balance.
27 balance_counter[balance] += 1
28
29 # Reset balance for the left side of k.
30 balance = 0
31
32 # Look at the elements to the left of k in nums.
33 for j in range(index_of_k - 1, -1, -1):
34 # Adjust balance.
35 balance += 1 if nums[j] > k else -1
36
37 # Count subarrays where balance is 0 or 1 just by k itself.
38 total_subarrays += 0 <= balance <= 1
39
40 # Use the counter balance to find valid subarrays that when combined
41 # with the current left part forms a balanced subarray.
42 total_subarrays += balance_counter[-balance] + balance_counter[-balance + 1]
43
44 # Return the total number of subarrays.
45 return total_subarrays
46
1class Solution {
2 public int countSubarrays(int[] nums, int k) {
3 // Getting the total number of elements in the nums array
4 int numberOfElements = nums.length;
5 // Iterator 'i' is used to find the position of the number 'k' in the array
6 int i = 0;
7 while (nums[i] != k) {
8 ++i;
9 }
10
11 // Initialize an array to keep track of counts
12 int[] counts = new int[numberOfElements * 2 + 1];
13 // Variable 'answer' is used for storing the final answer
14 int answer = 1;
15 // Variable 'x' stores the accumulated count
16 int x = 0;
17
18 // Check to the right of the position where k was found
19 for (int j = i + 1; j < numberOfElements; ++j) {
20 // If current element is greater than k, increment x; otherwise decrement x
21 x += nums[j] > k ? 1 : -1;
22 // If x is in the range of 0 to 1 (inclusive), increment the answer
23 if (x >= 0 && x <= 1) {
24 ++answer;
25 }
26 // Increment the count of the position x in the counts array
27 ++counts[x + numberOfElements];
28 }
29 // Reset 'x' for the second loop
30 x = 0;
31 // Check to the left of the position where k was found
32 for (int j = i - 1; j >= 0; --j) {
33 // If current element is greater than k, increment x; otherwise decrement x
34 x += nums[j] > k ? 1 : -1;
35 // If x is in the range of 0 to 1 (inclusive), increment the answer
36 if (x >= 0 && x <= 1) {
37 ++answer;
38 }
39 // Add to the answer the counts for -x and -x + 1 in the counts array
40 answer += counts[-x + numberOfElements] + counts[-x + 1 + numberOfElements];
41 }
42
43 // Returning the final answer which is the total count of valid sub-arrays
44 return answer;
45 }
46}
47
1#include <vector>
2#include <algorithm>
3#include <cstring>
4
5class Solution {
6public:
7 int countSubarrays(std::vector<int>& nums, int k) {
8 int numElements = nums.size();
9
10 // Find the position of the first occurrence of k.
11 int kPosition = std::find(nums.begin(), nums.end(), k) - nums.begin();
12
13 // Initialize a count array with double the size of nums and an extra one, and set to all zeros.
14 int countArray[(numElements << 1) | 1]; // Using bit manipulation to calculate 2*numElements+1 for size.
15 std::memset(countArray, 0, sizeof(countArray));
16
17 // Initialize answer with 1 counting the subarray that consists of only element k itself.
18 int answer = 1;
19
20 // This variable will store the relative comparison count.
21 int compCount = 0;
22
23 // Iterate through the elements to the right of k and count eligible subarrays.
24 for (int rightIndex = kPosition + 1; rightIndex < numElements; ++rightIndex) {
25 compCount += nums[rightIndex] > k ? 1 : -1;
26 if (compCount >= 0 && compCount <= 1) {
27 ++answer;
28 }
29 ++countArray[compCount + numElements];
30 }
31
32 // Reset comparison count for left side iteration.
33 compCount = 0;
34
35 // Iterate through the elements to the left of k and count eligible subarrays.
36 for (int leftIndex = kPosition - 1; leftIndex >= 0; --leftIndex) {
37 compCount += nums[leftIndex] > k ? 1 : -1;
38 if (compCount >= 0 && compCount <= 1) {
39 ++answer;
40 }
41 // Add count of subarrays matching the required condition.
42 answer += countArray[-compCount + numElements] + countArray[-compCount + 1 + numElements];
43 }
44
45 // Return the total count of subarrays.
46 return answer;
47 }
48};
49
1function countSubarrays(nums: number[], k: number): number {
2 // Find the index of value 'k' in the array.
3 const indexOfK = nums.indexOf(k);
4 const lengthOfNums = nums.length;
5
6 // Initialize a counting array to store frequencies of different sums.
7 // The size is made to be twice the length of the input array to avoid negative indexing.
8 // An extra slot is added to accommodate a '0' sum at the center.
9 const count = new Array((lengthOfNums << 1) | 1).fill(0);
10
11 // Initialize the answer with 1 to account for the 'k' itself.
12 let answer = 1;
13 let sum = 0;
14
15 // Count for subarrays starting from the index after 'k'.
16 for (let j = indexOfK + 1; j < lengthOfNums; ++j) {
17 // Ternary operator: if the current element is greater than 'k', increment sum,
18 // otherwise decrement sum.
19 sum += nums[j] > k ? 1 : -1;
20 // If sum is in [0,1] it is a valid subarray ending with nums[j].
21 answer += sum >= 0 && sum <= 1 ? 1 : 0;
22 // Increment the count of the current sum.
23 ++count[sum + lengthOfNums];
24 }
25
26 // Reset sum for counting subarrays starting before 'k'.
27 sum = 0;
28
29 // Count for subarrays starting before 'k'.
30 for (let j = indexOfK - 1; j >= 0; --j) {
31 // Increase or decrease sum depending on whether the value is greater than 'k'.
32 sum += nums[j] > k ? 1 : -1;
33 // If sum is in [0,1] it is a valid subarray starting with nums[j].
34 answer += sum >= 0 && sum <= 1 ? 1 : 0;
35 // Add the count of subarrays that would form a valid sum when combined.
36 answer += count[-sum + lengthOfNums] + count[-sum + 1 + lengthOfNums];
37 }
38
39 // Return the total count of valid subarrays.
40 return answer;
41}
42
Time and Space Complexity
Time Complexity
The time complexity of the given code can be analyzed in the following steps:
-
Finding the index of
k
innums
usingnums.index(k)
– This operation isO(n)
wheren
is the number of elements innums
. -
Iterating over the elements to the right of
k
, incrementingx
, and updating theans
and theCounter
for each element – The iteration is done once for each element to the right ofk
, which isO(n)
in the worst case (whenk
is the first element). -
Iterating over the elements to the left of
k
, incrementingx
, and updatingans
based on the currentx
value and the count from theCounter
– Similar to the step above, this is alsoO(n)
in the worst case (whenk
is the last element).
Given that these steps are performed sequentially, we add the complexities resulting in a total time complexity of O(n) + O(n) + O(n) = O(n)
.
Space Complexity
The space complexity of the code can be analyzed by considering the additional data structures used:
cnt
- ACounter
object which in the worst case stores counts for each possible distinct value ofx
. Sincex
represents a running difference between the count of values greater thank
and less thank
,x
can be at mostn
in absolute value in the worst case, leading to2n + 1
different possible values (including zero). Thus, theCounter
can have at mostO(n)
elements.
Hence, the space complexity of the code is O(n)
.
Learn more about how to find time and space complexity quickly using problem constraints.
What's the output of running the following function using input 56
?
1KEYBOARD = {
2 '2': 'abc',
3 '3': 'def',
4 '4': 'ghi',
5 '5': 'jkl',
6 '6': 'mno',
7 '7': 'pqrs',
8 '8': 'tuv',
9 '9': 'wxyz',
10}
11
12def letter_combinations_of_phone_number(digits):
13 def dfs(path, res):
14 if len(path) == len(digits):
15 res.append(''.join(path))
16 return
17
18 next_number = digits[len(path)]
19 for letter in KEYBOARD[next_number]:
20 path.append(letter)
21 dfs(path, res)
22 path.pop()
23
24 res = []
25 dfs([], res)
26 return res
27
1private static final Map<Character, char[]> KEYBOARD = Map.of(
2 '2', "abc".toCharArray(),
3 '3', "def".toCharArray(),
4 '4', "ghi".toCharArray(),
5 '5', "jkl".toCharArray(),
6 '6', "mno".toCharArray(),
7 '7', "pqrs".toCharArray(),
8 '8', "tuv".toCharArray(),
9 '9', "wxyz".toCharArray()
10);
11
12public static List<String> letterCombinationsOfPhoneNumber(String digits) {
13 List<String> res = new ArrayList<>();
14 dfs(new StringBuilder(), res, digits.toCharArray());
15 return res;
16}
17
18private static void dfs(StringBuilder path, List<String> res, char[] digits) {
19 if (path.length() == digits.length) {
20 res.add(path.toString());
21 return;
22 }
23 char next_digit = digits[path.length()];
24 for (char letter : KEYBOARD.get(next_digit)) {
25 path.append(letter);
26 dfs(path, res, digits);
27 path.deleteCharAt(path.length() - 1);
28 }
29}
30
1const KEYBOARD = {
2 '2': 'abc',
3 '3': 'def',
4 '4': 'ghi',
5 '5': 'jkl',
6 '6': 'mno',
7 '7': 'pqrs',
8 '8': 'tuv',
9 '9': 'wxyz',
10}
11
12function letter_combinations_of_phone_number(digits) {
13 let res = [];
14 dfs(digits, [], res);
15 return res;
16}
17
18function dfs(digits, path, res) {
19 if (path.length === digits.length) {
20 res.push(path.join(''));
21 return;
22 }
23 let next_number = digits.charAt(path.length);
24 for (let letter of KEYBOARD[next_number]) {
25 path.push(letter);
26 dfs(digits, path, res);
27 path.pop();
28 }
29}
30
Recommended Readings
Prefix Sum The prefix sum is an incredibly powerful and straightforward technique Its primary goal is to allow for constant time range sum queries on an array What is Prefix Sum The prefix sum of an array at index i is the sum of all numbers from index 0 to i By
LeetCode Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way we
Recursion Recursion is one of the most important concepts in computer science Simply speaking recursion is the process of a function calling itself Using a real life analogy imagine a scenario where you invite your friends to lunch https algomonster s3 us east 2 amazonaws com recursion jpg You first