Leetcoding in Rust—Curious case of Sub Arrays and its relationship with K
Getting started with Graphs, Trees, and BSTs is harder, but it gets easier with practice. On the other hand, getting started with Arrays and LinkedLists is always easy, but it gets harder as you do more problems.
- By an average Leetcoder (Failed FAANG Interviewee)
Love stories of the Sub Array and K can be found all over Leetcode and is nothing new to an average Leetcoder. But the more you think you are done with their story, the more they pull you into their lore.
It all started with a naive Subarray in the Medium section of Leetcode in search of some K. With practice, you know how it all ends, but do you?
- Subarray/Substring — Use Sliding Window
- Subsequence — Use Dynamic Programming
- Subsets — Use Backtracking
And you merrily go your way and live happily ever after. You wish it ended like that as you turn the page to face the Hards. You are greeted with a love affair with all the ingredients for making a whole DS&A chapter.
Chapter 1 — The Sliding Window
992. Subarrays with K Different Integers
This problem uses a Sliding Window but with an exciting twist. But before that, we need to understand the attractive property of counting subarrays.
Let's find the subarrays where the number of distinct elements (integers) is Atmost K (let's take the value of K = 3). How can we do that?
Input = [1, 2, 3], K = 3
Then if we iterate over all possible solutions, it will be:
[1],
[2],
[3],
[1, 2],
[2, 3],
[1, 2, 3]
So, the count is 6.
But how can we count this if this K was large and Input size was arbitarylarge?Here is the sudo code for how to count subarrays from 0 to n elements with some constraints (let's call that invalid_state()):
let inpt = vec![1, 2, 3]
let n = inpt.len() # 3
let (mut i, mut j, mut count) = (0usize, 0usize, 0usize)
fn<T> modify_state<T>(some_state: T, index: i32) {
....
}
fn<T> invalid_state<T>(some_state: T) -> bool {
true
}
while j < n {
modify_state(j, some_state);
while i < j && invalid_state(some_state) { // When constraints fail
modify_state(i, some_state);
i += 1;
}
count += j - i + 1;
j += 1;
}
return countSo, we arrive at a general solution to count the number of subarrays with almost every size that meets some constraints. You can put anything in here from [(1) Atmost Distinct K Elements, (2) Elements divisible by 3, (3) Even integers…] in place of invalid_state. In our case, we will check if the number of distinct elements is almost K.
The other concept is counting Exactly K Elements (Not Atmost K). There are several ways to do this:
- Total Subarrays minus Total subarrays that do not contain Exactly K elements.
- Total Subarrays contain Atmost K Distinct Elements minus Total Subarrays that contain Atmost K — 1 Distinct Elements.
- Total Subarrays contain Atleast K+1 Distinct Elements minus Total Subarrays that contain Atleast K Distinct Elements.
We can leverage any of the above ways to find Exactly K. But as we already know how to find Atmost K, we can use (2) to solve the problem. You should be able to solve it now on your own, but here is the solution.
use std::collections::HashMap;
impl Solution {
fn subarrays_with_atmost_k_distinct(nums: &Vec<i32>, k: i32) -> i32 {
if k == 0 {
return 0;
}
let (mut i, mut j, mut res, mut hashmap) = (0usize, 0usize, 0usize, HashMap::<i32, i32>::new());
while j < nums.len() {
let key_j = nums[j];
if !hashmap.contains_key(&key_j) {
hashmap.insert(key_j, 0);
}
hashmap.insert(key_j, *hashmap.get(&key_j).unwrap() + 1);
while i < j && hashmap.len() > (k as usize) {
let key_i = nums[i];
hashmap.insert(key_i, *hashmap.get(&key_i).unwrap() - 1);
if *hashmap.get(&key_i).unwrap() == 0 {
hashmap.remove(&key_i);
}
i += 1;
}
res += j - i + 1;
j += 1;
}
res as i32
}
pub fn subarrays_with_k_distinct(nums: Vec<i32>, k: i32) -> i32 {
return Self::subarrays_with_atmost_k_distinct(&nums, k) - Self::subarrays_with_atmost_k_distinct(&nums, k - 1);
}
}Chapter 2 — The Tale of Two Pointers
1793. Maximum Score of a Good Subarray
This problem looks like a simple subarray problem with k elements, but it's nothing like that. What the problem is asking you to do is include the kth element in the final subarray.
So you start from the kth element and compare the elements on the left and right to see if they change the min element, and if it does, then does it produce the maximum value? It's a Two-pointers problem, where you extend the left and right indexes based on the criteria.
Interestingly, the problem is that (j — i + 1) can never go down as you can increase j to the right or decrease i to the left. So, what matters in the problem is the minimum elements to the left and right. So, if the left element is greater than the element on the right, we can try to expand on the left and see if we get a better result. Otherwise, we can expand on the right to see if we get a better result. This is a Two Pointers + Greedy problem.
use std::cmp;
impl Solution {
pub fn maximum_score(nums: Vec<i32>, k: i32) -> i32 {
let (mut i, mut j, mut curMin, mut res) = (k as usize, k as usize, nums[k as usize] as usize, nums[k as usize] as usize);
while j < nums.len() - 1 || i > 0 {
let left = if (i > 0) { nums[i - 1] as usize } else { 0 as usize };
let right = if (j < nums.len() - 1) { nums[j + 1] as usize } else { 0 as usize };
if left > right {
i -= 1;
curMin = cmp::min(curMin, left);
} else {
j += 1;
curMin = cmp::min(curMin, right)
}
res = cmp::max(res, curMin * (j - i + 1))
}
res as i32
}
}Chapter 3 — The Order of the Deque
862. Shortest Subarray with Sum at Least K
This problem resembles the Shortest Subarray with Sum K (or Exactly K), which a Sliding Window can easily solve. However, this problem is very different as we must find the Shortest Subarray with Sum K or above. This means if K is 10, and the Shortest Subarray we get with Sum 10 is 20, while we have another Subarray with Sum 11 but the size of the Subarray is 19, then the answer will be 19.
The Brute force way of doing this would be to get all Subarrays, filter down to Subarrays whose Sum is K or above, and then find the Subarray with the minimum size. But this will be O(n²). We need to do better as the maximum length is of order 10⁵.
One of the optimal approaches is to use a Prefix Sum array to calculate the Subarray Sums. How does that work? Let me make that clear here:
Input = [2,-1,2] # Given Subarray
K = 3
PrefixSums = [0, 2, 1, 3] # We calculate the Prefix Sums (Cumulative SSumSums)After we have the Prefix Sums, we can easily calculate the Subarrays. For example, let's say we want to know the Sum of the Subarray from Index = 1 to Index = 2.
sumOfSubarrayIndexOneToTwo = PrefixSums[2 + 1] - PrefixSums[1]
= PrefixSums[3] - PrefixSums[1]
= 3 - 2
= 1
We can check manually that the sum of subarray [-1, 2] is indeed equal to 1. Now that we know how to calculate the Sum of Subarrays from Prefix Sum, we need to find the Index till we need to stretch the Subarray for each Index till the Sum is equal to or greater than K. Basically, we need to find the two indexes between which the Sum is equal to or greater than K.
For this, we have this equation:
PrefixSum[end + 1] - PrefixSum[start] >= KWe can rewrite this as:
PrefixSum[end + 1] - K >= PrefixSum[start]So we can iterate through the Prefix Sums by each Index and then iterate again till that Index to see which start value satisfies the above equation. Then the size of the Subarray will be end + 1 — start. Great, but now the problem is we are still making two loops, so the Runtime will still be O(n²); it doesn’t help.
But this is where we need the power of a data structure that can keep track of values we care about and discard those we don’t so that we don’t end up iterating over them again.
Let’s try to visualize this:
Input = [2,-1,2] # Given Subarray
K = 3
PrefixSums = [0, 2, 1, 3] # We calculate the Prefix Sums (Cumulative Sums)
So for each Index,
let's apply the above equation (PrefixSums[end + 1] - K >= PrefixSums[start]):
1. Index = 0, End = 0:
--------------------------
PrefixSums[0 + 1] - K
=> 2 - 3
=> -1
Is there a value in the PrefixSums till Index = 0, that has the value -1 or less? No.
So nothing to do. Let's move on.
2. Index = 1, End = 1:
--------------------------
PrefixSums[1 + 1] - K
=> 1 - 3
=> -2
Is there a value in the PrefixSums till Index = 1, that has the value -2 or less? No.
So nothing to do. Let's move on.
2. Index = 1, End = 1:
--------------------------
PrefixSums[2 + 1] - K
=> 3 - 3
=> 0
Is there a value in the PrefixSums till Index = 2, that has the value 0 or less? Yes.
The value at index = 0, has the value = 0 in the PrefixSums.
So the size of the Subarray that satisfies the condition is 2 + 1 - 0 = 3
And 3 is the answer.
In this example, we didn’t need to iterate over every value in the PrefixSum array for each Index as the first element had satisfied the condition. Still, let’s take the following example (Taken from the explanation here):
Input = [3, 5, 2, 1, -6, 7, 1, 8]
K = 8
PrefixSums = [0, 3, 8, 10, 11, 5, 12, 13, 21]
1. Index = 0, End = 0:
--------------------------
PrefixSums[0 + 1] - K
=> 3 - 8
=> -5
So do we have -5 or less in the PrefixSums? No. We move on.
2. Index = 1, End = 1:
--------------------------
PrefixSums[1 + 1] - K
=> 8 - 8
=> 0
So do we have -5 or less in the PrefixSums? Yes.
We record the size = 1 + 1 - 0 = 2
3. Index = 2, End = 2:
--------------------------
PrefixSums[2 + 1] - K
=> 10 - 8
=> 2
So do we have 2 or less in the PrefixSums? Yes, its 0.
But hold on we have already used 0 and if we use it again the size of the
subarray is going to be larger than what we have found earlier.
Size = 2 + 1 - 0 = 3.
So that doesn't help.
Why not just discard this value when we get a match as another match will not
give us a smaller size subarray. Remember we need to find the Shortest Subarray.This is very interesting. If we find a matching pair, anything after it will not add any value if we take the same start Index in the PrefixSum, as it will be bigger. Ok, so let’s remove it from the PrefixSums.
2. Index = 1, End = 1:
--------------------------
PrefixSums[1 + 1] - K
=> 8 - 8
=> 0
So do we have -5 or less in the PrefixSums? Yes.
We record the size = 1 + 1 - 0 = 2.
Lets remove the element at Index = 0.
The new PrefixSums subarray:
PrefixSums = [N/A, 3, 8, 10, 11, 5, 12, 13, 21]. Ok lets continue.
3. Index = 2, End = 2:
--------------------------
PrefixSums[2 + 1] - K
=> 10 - 8
=> 2
So do we have 2 or less in the PrefixSums? No. We move on.
4. Index = 3, End = 3:
--------------------------
PrefixSums[3 + 1] - K
=> 11 - 8
=> 3
So do we have 3 or less in the PrefixSums? Yes.
We record the size = 3 + 1 - 1 = 3. But last time we got 2, so this is not the
shortest. So we can discard it.
Lets remove the element at Index = 1.
The new PrefixSums subarray:
PrefixSums = [N/A, N/A, 8, 10, 11, 5, 12, 13, 21]. Ok lets continue.
5. Index = 4, End = 4:
--------------------------
PrefixSums[4 + 1] - K
=> 5 - 8
=> -3
So do we have 3 or less in the PrefixSums? No. We move on.
6. Index = 5, End = 5
--------------------------
PrefixSums[5 + 1] - K
=> 12 - 8
=> 4
So do we have 4 or less in the PrefixSums? No. We move on.
7. Index = 6, End = 6
--------------------------
PrefixSums[6 + 1] - K
=> 13 - 8
=> 5
So do we have 5 or less in the PrefixSums? Yes.
We record the size = 6 + 1 - 5 = 2. But this same as 2.
Lets remove the element at Index = 5.
The new PrefixSums subarray:
PrefixSums = [N/A, N/A, 8, 10, 11, N/A, 12, 13, 21]. Ok lets continue.
But something interesting has happened now, the elements before 5:
8, 10, 11 got skipped.
Do we need to keep them in here anymore?
Even if we do get 8 or 10 or 11 the size of the subarray is going to be bigger
than 2. So what is the point of keeping them?
We can also just nullify them.
PrefixSums = [N/A, N/A, N/A, N/A, N/A, N/A, 12, 13, 21]So we see two separate ways we need to nullify the elements: whatever comes to the left when it matches, or else if elements got skipped, we might have to go back from the right and nullify the elements.
This is where we see the power of the Deque. Instead of keeping the PrefixSums in an array, what if we add them in a Deque? Then, we don’t need to reiterate the nullified values; we can remove them with O(1) operation.
Also, we don’t need to precalculate the entire PrefixSums; we only need the PrefixSum at the end + 1 and whatever comes before. So, we can keep a single variable to store the PrefixSum of end + 1 and then add it to our Deque after we compare the values already in the Deque.
So, now that we understand the solution let's code it up with a Deque.
impl Solution {
pub fn shortest_subarray(nums: Vec<i32>, k: i32) -> i64 {
let mut deq: std::collections::VecDeque<(i64, i64)> = std::collections::VecDeque::new();
deq.push_back((0, -1));
let (mut prefixSum, mut res) = (0 as i64, i64::MAX);
for (idx, val) in nums.iter().enumerate() {
prefixSum += (*val as i64);
while deq.len() > 0 && prefixSum <= deq.back().unwrap().0 {
deq.pop_back();
}
while deq.len() > 0 && prefixSum - (k as i64) >= deq.front().unwrap().0 {
res = std::cmp::min(res, (idx as i64) - deq.front().unwrap().1);
deq.pop_front();
}
deq.push_back((prefixSum, idx as i64))
}
if res == i64::MAX {
return -1;
}
res
}
}Chapter 4 — The Unusual Sliding Window
2444. Count Subarrays With Fixed Bounds
This problem can be solved by the Sliding Window technique but with a very different twist. This will require us to keep track of 4 different indexes, unlike a usual Sliding Window problem where we keep track of 2 indexes.
If you think about the problem, you must keep track of a few states, i.e., (1) Where does the window start? (2) Where does it end? (3) Have we found the element with the minimum value, which is given as minK, and (4) have we found the element with the maximum value, which is given as maxK?
Let's try first to solve the problem where all elements are between minK and maxK. Then, try to extend the solution to the actual problem.
input = [3,2,1,2,1,5,3]
minK = 1, maxK = 5
So we can see all elements are between minK (1) and maxK (5).
So how many valid subarrays possible?
Subarrays:
[3], [2], [1], [2], [1], [5], [3] -> X Not valid
[3, 2], [2, 1], [1, 2], [2, 1], [5, 3] -> X Not valid
[1, 5] -> Valid!
[3, 2, 1], [2, 1, 2], [1, 2, 1] -> X Not valid
[2, 1, 5], [1, 5, 3] -> Valid!
[3, 2, 1, 2] -> X Not valid
[2, 1, 2, 1] -> X Not valid
[1, 2, 1, 5], [2, 1, 5, 3] -> Valid!
[3, 2, 1, 2, 1] -> X Not valid
[2, 1, 2, 1, 5], [1, 2, 1, 5, 3] -> Valid!
[3, 2, 1, 2, 1, 5], [2, 1, 2, 1, 5, 3] -> Valid!
[3, 2, 1, 2, 1, 5, 3] Valid!
So there are only 10 solutions here.
One thing to notice is if we see the solutions and see it from the right:
[1, 5]
[2, 1, 5]
[1, 2, 1, 5]
[2, 1, 2, 1, 5]
[3, 2, 1, 2, 1, 5]
So if we see from the right ones we have both the elements minK and maxK,
then we can extend it till by any number of element and it will be valid.
[1, 5, 3]
[2, 1, 5, 3]
[1, 2, 1, 5, 3]
[2, 1, 2, 1, 5, 3]
[3, 2, 1, 2, 1, 5, 3]
Similarly, if we have both minK and maxK, we can also extend it to the right,
and extend it as much as possible.
So this can be thought of as a sliding window, where we take i, j as the
start and end of the dynamically sized sliding window, where we move j from
left to right, keep track of the number of times we get minK and maxK. Then we shrink
the sliding window, i.e. move i to the right till we have both 1 or more minK
and maxK element tracked. Then we can add i + 1 to the result. We do this
everytime we move j to the right and we still have 1 or more minK and maxK
tracked.
We take a hashmap = {} to track minK and maxK. We then take 2 pointers to be
the 2 ends of sliding window i, j.
[3, 2, 1, 2, 1, 5, 3] {}
^^
ij
Then,
[3, 2, 1, 2, 1, 5, 3] {}
^ ^
i j
Then,
[3, 2, 1, 2, 1, 5, 3] {1: 1} # We have found minK (1) and 1 times
^ ^
i j
Then,
[3, 2, 1, 2, 1, 5, 3] {1: 1}
^ ^
i j
Then,
[3, 2, 1, 2, 1, 5, 3] {1: 2} # We found another minK
^ ^
i j
Then,
[3, 2, 1, 2, 1, 5, 3] {1: 2, 5: 1} # Now we have both 1 and 5
^ ^
i j
Now we start moving i,
[3, 2, 1, 2, 1, 5, 3] {1: 2, 5: 1}
^ ^
i j
Then,
[3, 2, 1, 2, 1, 5, 3] {1: 2, 5: 1}
^ ^
i j
Now we remove 1 minK, and keep moving
[3, 2, 1, 2, 1, 5, 3] {1: 1, 5: 1}
^ ^
i j
Then,
[3, 2, 1, 2, 1, 5, 3] {1: 1, 5: 1}
^ ^
i j
Now if we move i one more time to the right we will no longer have
a valid subarray so we stop and add i + 1 to the result
result += i + 1 = 4 + 1 = 5
Now we can move j again,
[3, 2, 1, 2, 1, 5, 3] {1: 1, 5: 1}
^ ^
i j
Now we check again and find we meet the criteria, and try to move i,
but moving i will again lead to invalid subarray, we again collect the
results with adding i + 1 to the result
result += i + 1 = 5 + 4 + 1 = 10
And finally we have found the result.Even though we have found the result, this is still O(n²). But we can do another interesting optimization if we understand an interesting property of the iteration. We can keep track of the last minK and maxK position instead of the count, and every time we have both minK and maxK position, we add the minimum of the position of minK and maxK. It's the same thing, but we can eliminate the i-pointer iteration.
[3, 2, 1, 2, 1, 5, 3] minK_pos = -1, maxK_pos = -1 # We track the positions of minK and maxK
^
j
Then,
[3, 2, 1, 2, 1, 5, 3] minK_pos = -1, maxK_pos = -1
^
j
Now we found the first minK position,
[3, 2, 1, 2, 1, 5, 3] minK_pos = 2, maxK_pos = -1
^
j
Then,
[3, 2, 1, 2, 1, 5, 3] minK_pos = 2, maxK_pos = -1
^
j
We found the second/last minK position,
[3, 2, 1, 2, 1, 5, 3] minK_pos = 4, maxK_pos = -1
^
j
Next we also found the maxK position, now we can add min(minK_pos, maxK_pos) + 1
[3, 2, 1, 2, 1, 5, 3] minK_pos = 4, maxK_pos = 5
^
j
Now we add min(minK_pos, maxK_pos) + 1 = 4 + 1,
result += 5
Then, we continue
[3, 2, 1, 2, 1, 5, 3] minK_pos = 4, maxK_pos = 5
^
j
We still have both the same positions, so we can again add the same
result += min(minK_pos, maxK_pos) + 1 = 5 + 4 + 1 = 10Awesome, we have found an algorithm to find all subarrays between minK and maxK if all elements of the original array are between minK and maxK. We now need the final piece, that is, when we have elements not within minK and maxK.
This is very easy to see,
Lets say this is the new input array: [3, 2, 1, 2, 7, 1, 5, 3]
Here we can see 7 which wasn't there before, and it does not fall within minK (1) and maxK (5)
So we can never include this 7 in our subarrays, and we can see this as a partition.
Our answers are going to be within the 2 splits:
[3, 2, 1, 2] and [1, 5, 3]
And instead of adding, min(minK_pos, maxK_pos) + 1 to the result, we need to add
min(minK_pos, maxK_pos) - invalid_pos
where invalid_pos is the last index of element that splits the array
Also remember when we split the array the minK_pos and maxK_pos has to be reset.
Lets try the new algorithm for the input: [3, 2, 1, 2, 7, 1, 5, 3]
[3, 2, 1, 2, 7, 1, 5, 3] minK_pos = -1, maxK_pos = -1, invalid_pos = -1
^
j
Then,
[3, 2, 1, 2, 7, 1, 5, 3] minK_pos = -1, maxK_pos = -1, invalid_pos = -1
^
j
We found the first minK pos,
[3, 2, 1, 2, 7, 1, 5, 3] minK_pos = 2, maxK_pos = -1, invalid_pos = -1
^
j
Then,
[3, 2, 1, 2, 7, 1, 5, 3] minK_pos = 2, maxK_pos = -1, invalid_pos = -1
^
j
Now we hit the invalid number, so we track that and reset, basically split
[3, 2, 1, 2, 7, 1, 5, 3] minK_pos = -1, maxK_pos = -1, invalid_pos = 4
^
j
We found the first minK pos,
[3, 2, 1, 2, 7, 1, 5, 3] minK_pos = 5, maxK_pos = -1, invalid_pos = 4
^
j
Now we found the maxK pos too,
[3, 2, 1, 2, 7, 1, 5, 3] minK_pos = 5, maxK_pos = 6, invalid_pos = 4
^
j
So we now add min(minK_pos, maxK_pos) - invalid_pos to result,
result += min(minK_pos, maxK_pos) - invalid_pos = min(5, 6) - 4 = 5 - 4 = 1
Then,
[3, 2, 1, 2, 7, 1, 5, 3] minK_pos = 5, maxK_pos = 6, invalid_pos = 4
^
j
So nothing changed, and we still have minK_pos and maxK_pos found and not reset,
so we add again min(minK_pos, maxK_pos) - invalid_pos to the result.
result += min(minK_pos, maxK_pos) - invalid_pos = 1 + 1 = 2
So now we result 2. And we can also manually verify:
[1, 5]
[1, 5, 3]
are the 2 valid subarrays.Now, here is the implementation,
use std::cmp::{min, max};
impl Solution {
pub fn count_subarrays(nums: Vec<i32>, minK: i32, maxK: i32) -> i64 {
let (mut j, mut minK_pos, mut maxK_pos, mut invalid_pos, mut res) = (0usize, -1i64, -1i64, -1i64, 0i64);
while j < nums.len() {
if nums[j] < minK || nums[j] > maxK {
minK_pos = -1i64;
maxK_pos = -1i64;
invalid_pos = j as i64;
} else {
if nums[j] == minK {
minK_pos = j as i64;
}
if nums[j] == maxK {
maxK_pos = j as i64;
}
if minK_pos != -1 && maxK_pos != -1 {
res += min(minK_pos, maxK_pos) - invalid_pos;
}
}
j += 1;
}
res
}
}Chapter 5 — What only matters is whether it is greater or lesser than K.
2488. Count Subarrays With Median K
To solve this problem, we need several concepts, just like all the problems above, which I would go over. Most hard problems are hard because they involve multiple steps. This is true here, too.
The first concept that we need to understand is the median element, which is the element that comes just in the middle for odd-sized arrays, and in this problem, for even-sized arrays, we are supposed to take a left-middle element. So what matters is to see if K is in the middle or left-middle of our subarray; if it is, then it is a valid subarray, which we need to consider.
So, we can change the input into elements that are lesser than K with -1, K with 0, and then elements greater than K with 1. We will see why this change makes sense in the 2nd concept; let's do that for now.
input = [3,2,1,4,5]
K = 4
After the transformation:
input = [-1, -1, -1, 0, 1]
We can see that our subarrays are going to be:
[0]
[0, 1]
[-1, 0, 1]Great, now we can also see that if we add up the elements of the valid subarrays, it comes between [0, 1]. More precisely, the sum is 0 for an odd-sized subarray, while for an even-sized subarray, it's 1.
But before we proceed further, let us solve another problem: Count all subarrays that sum up to 0 from the given array.
input = [-1, -1, -1, 0, 1]
To solve this, we can use prefix sum array. Which is to calculate the cumulative
sums by elements.
prefix_sum_array = [0, -1, -2, -3, -3, -2]
The way to calculate the prefix_sum_array is to start with 0, then add:
prefix_sum_array[i + 1] = prefix_sum_array[i] + input[i]
Now to calculate the sum of a subarray i to j, we need to use:
sum_of_subarray_i_to_j = prefix_sum_array[j + 1] - prefix_sum_array[i]
So lets say we want to find the sum of subarray of elements from index 2 to 4, we would:
sum_of_subarray_2_to_4 = prefix_sum_array[4 + 1] - prefix_sum_array[2]
= -2 - (-2)
= 0
So we can check with the input array that indeed its 0: sum([-1, 0, 1])
So now that we understand the concept of prefix_sum_arrays, how do we use this to solve the new problem: Count all subarrays that sum up to 0 from the given array?
/**
To solve this problem what we need to so is for some index = j, we need to find
some index = i that when substracted gives us 0 as the result. More precisely,
we need to find some i which has the same value as the index j.
Here is the sudo code on how to do this:
**/
fn countOfSubarraysWithSumZero(input: Vec<i32>) {
let mut prefix_sum_arrays = [0; input.len() + 1];
let (mut j, mut prev_hashmap, mut res) = (1usize, HashMap::<usize, usize>::new(), 0usize);
prev_hashmap.insert(0, 1);
while j < prefix_sum_arrays.len() {
let key_j = prefix_sum_arrays[j] as usize;
if prev_hashmap.contains_key(&key_j) {
res += *prev_hashmap.get(&key_j).unwrap();
} else {
prev_hashmap.insert(key_j, 0);
}
prev_hashmap.insert(key_j, *prev_hashmap.get(&key_j).unwrap() + 1);
}
res
}
So, we have solved this, and isn’t this similar to the problem we have above the original problem? Find all subarrays that have sum = 0 or 1. So, let's continue the original problem.
input = [3,2,1,4,5]
K = 4
After the transformation:
input = [-1, -1, -1, 0, 1]
We can see that our subarrays are going to be:
[0]
[0, 1]
[-1, 0, 1]
Lets now calculate the prefix_sum_arrays:
prefix_sum_arrays = [0, -1, -2, -3, -3, -2]
Now, lets try to find all the subarrays that sum up to 0 or 1:
result = 0
-----------------------------------------
j = 1
hashmap = {0: 1}
current_sum = prefix_sum_arrays[1] = -1
To sum up to 0, we need to find -1 because
prefix_sum_array[j + 1] - prefix_sum_array[i] = 0
=> prefix_sum_array[i] = prefix_sum_array[j + 1]
And To sum up to 1, we need to find -2 because
prefix_sum_array[j + 1] - prefix_sum_array[i] = 1
=> prefix_sum_array[i] = prefix_sum_array[j + 1] - 1
We did not find either, so lets just add the -1 in the hashmap and continue.
result = 0
-----------------------------------------
j = 2
hashmap = {0: 1, -1: 1}
current_sum = prefix_sum_arrays[2] = -2
We need -2 or -3 in the hashmap, we didn't find it in the hashmap so we add the
current element -2 and move on.
result = 0
-----------------------------------------
j = 3
hashmap = {0: 1, -1: 1, -2: 1}
current_sum = prefix_sum_arrays[3] = -3
We need -3 or -4 in the hashmap, we didn't find it in the hashmap so we add the
current element -3 and move on.
result = 0
-----------------------------------------
j = 4
hashmap = {0: 1, -1: 1, -2: 1, -3: 1}
current_sum = prefix_sum_arrays[4] = -3
We need -3 or -4 in the hashmap, we finally found -3 in the hashmap,
we add the number of times we found -3 which 1 to the result.
result = 1
This is where things differ a bit, we don't add -3 again, as we have found a match.
-----------------------------------------
j = 5
hashmap = {0: 1, -1: 1, -2: 1, -3: 1}
current_sum = prefix_sum_arrays[5] = -2
We need -2 or -3 in the hashmap, we have both in the hashmap, we add the
number of times we saw both.
result += hashmap[-2] + hashmap[-3]
result = 3
So finally our result is 3. Which is what we showed above, the subarrays will:
[0]
[0, 1]
[-1, 0, 1]
Which translates to the original array as:
[4]
[4, 5]
[1, 4, 5]Now, there is a small caveat in this solution: there can be cases where we get the same element or we element — 1 in the hashmap, but we did not include any element with K as the value. In our current example, by chance, we got the value 0 in the transformed array, k=4 in the original array, when we found j = 4 and j = 5. So how can we avoid this situation?
We can only start considering matches if the index (our j) is at the index where we have the element = K or beyond that index. Here is the pseudo-code with the change:
input = [3,2,1,4,5]
K = 4
After the transformation:
input = [-1, -1, -1, 0, 1]
We can see that our subarrays are going to be:
[0]
[0, 1]
[-1, 0, 1]
Lets now calculate the prefix_sum_arrays:
prefix_sum_arrays = [0, -1, -2, -3, -3, -2]
Now, lets try to find all the subarrays that sum up to 0 or 1:
result = 0
-----------------------------------------
j = 1
hashmap = {0: 1}
current_sum = prefix_sum_arrays[1] = -1
Index is not equal to or greater than index = 4 in the prefix_sum_array,
so lets just add the -1 in the hashmap and continue.
result = 0
-----------------------------------------
j = 2
hashmap = {0: 1, -1: 1}
current_sum = prefix_sum_arrays[2] = -2
Index is not equal to or greater than index = 4 in the prefix_sum_array,
so lets just add the -2 in the hashmap and continue.
result = 0
-----------------------------------------
j = 3
hashmap = {0: 1, -1: 1, -2: 1}
current_sum = prefix_sum_arrays[3] = -3
Index is not equal to or greater than index = 4 in the prefix_sum_array,
so lets just add the -2 in the hashmap and continue.
result = 0
-----------------------------------------
j = 4
hashmap = {0: 1, -1: 1, -2: 1, -3: 1}
current_sum = prefix_sum_arrays[4] = -3
Now we are at index = 4, so we can compare.
We need -3 or -4 in the hashmap, we finally found -3 in the hashmap,
we add the number of times we found -3 which 1 to the result.
result = 1
This is where things differ a bit, we don't add -3 again, as we have found a match.
-----------------------------------------
j = 5
hashmap = {0: 1, -1: 1, -2: 1, -3: 1}
current_sum = prefix_sum_arrays[5] = -2
We need -2 or -3 in the hashmap, we have both in the hashmap, we add the
number of times we saw both.
result += hashmap[-2] + hashmap[-3]
result = 3
So finally our result is 3. Which is what we showed above, the subarrays will:
[0]
[0, 1]
[-1, 0, 1]
Which translates to the original array as:
[4]
[4, 5]
[1, 4, 5]Now that will work! Let's get it coded.
use std::collections::HashMap;
impl Solution {
pub fn count_subarrays(nums: Vec<i32>, k: i32) -> i32 {
let n = nums.len();
let mut new_nums = vec![0];
new_nums.resize(n, 0);
for i in 0..n {
if nums[i] < k {
new_nums[i] = -1;
}
else if nums[i] > k {
new_nums[i] = 1;
} else {
new_nums[i] = 0;
}
}
let mut prefix_sum_arrays = vec![0];
prefix_sum_arrays.resize(n + 1, 0);
let mut index_where_zero = 0;
for i in 0..n {
prefix_sum_arrays[i + 1] = prefix_sum_arrays[i] + new_nums[i];
if new_nums[i] == 0 {
index_where_zero = i + 1;
}
}
let (mut j, mut hashmap, mut res) = (1, HashMap::<i32, i32>::new(), 0);
hashmap.insert(0, 1);
while j < prefix_sum_arrays.len() {
let temp = prefix_sum_arrays[j] - 1;
if j >= index_where_zero {
// Find 0
if hashmap.contains_key(&prefix_sum_arrays[j]) {
res += *hashmap.get(&prefix_sum_arrays[j]).unwrap();
}
// Find 1
if hashmap.contains_key(&temp) {
res += *hashmap.get(&temp).unwrap();
}
} else {
let key = prefix_sum_arrays[j];
if !hashmap.contains_key(&key) {
hashmap.insert(key, 0);
}
hashmap.insert(key, *hashmap.get(&key).unwrap() + 1);
}
j += 1;
}
res
}
}
