subarray sum problem, two solutions proposed

2023_10_12/subarray-sum-equals-k/Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "subarray-sum-equals-k"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]

2023_10_12/subarray-sum-equals-k/src/main.rs

@@ -0,0 +1,40 @@
+// 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
+}