You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
41 lines
1.3 KiB
Rust
41 lines
1.3 KiB
Rust
// This first solution is O(n) time and O(n) space
|
|
|
|
pub fn subarray_sum_1(nums: Vec<i32>, k: i32) -> i32 {
|
|
let mut count = 0; // count of subarray
|
|
let mut sum = 0; // sum of subarray
|
|
let mut map = std::collections::HashMap::new(); // sum -> count, like Counter in python
|
|
map.insert(0, 1); // init sum = 0, count = 1
|
|
for i in 0..nums.len() { // iterate nums
|
|
sum += nums[i]; // sum of subarray
|
|
if let Some(v) = map.get(&(sum - k)) { // if sum - k in map, add count
|
|
count += v;
|
|
}
|
|
*map.entry(sum).or_insert(0) += 1; // add sum to map or add count
|
|
}
|
|
count
|
|
}
|
|
|
|
// This second solution is O(n^2) time and O(n) space and uses prefix sums. It is slower than the first solution but more didactic and efficient memory-wise.
|
|
|
|
pub fn subarray_sum2(nums: Vec<i32>, k: i32) -> i32 {
|
|
let mut count = 0; // count of subarrays with sum equal to k
|
|
let prefix_sums = nums
|
|
.iter()
|
|
.scan(0, |sum, &element| {
|
|
*sum += element;
|
|
Some(*sum)
|
|
})
|
|
.collect::<Vec<i32>>();
|
|
|
|
for i in 0..prefix_sums.len() {
|
|
for j in i..prefix_sums.len() {
|
|
let sum = if i == 0 { prefix_sums[j] } else { prefix_sums[j] - prefix_sums[i - 1] };
|
|
if sum == k {
|
|
count += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
count
|
|
}
|