Choose Numbers From Two Arrays in Range

Problem

You are given two 0-indexed integer arrays nums1 and nums2 of length n.

A range [l, r] (inclusive) where 0 <= l <= r < n is balanced if:

For every i in the range [l, r], you pick either nums1[i] or nums2[i].
The sum of the numbers you pick from nums1 equals to the sum of the numbers you pick from nums2 (the sum is considered to be 0 if you pick no numbers from an array).

Two balanced ranges from [l1, r1] and [l2, r2] are considered to be different if at least one of the following is true:

l1 != l2
r1 != r2
nums1[i] is picked in the first range, and nums2[i] is picked in the second range or vice versa for at least one i.

Return _the number ofdifferent ranges that are balanced. _Since the answer may be very large, return it modulo 109 + 7 .

Examples

Example 1:

Input: nums1 = [1,2,5], nums2 = [2,6,3]
Output: 3
Explanation: The balanced ranges are:
- [0, 1] where we choose nums2[0], and nums1[1].
  The sum of the numbers chosen from nums1 equals the sum of the numbers chosen from nums2: 2 = 2.
- [0, 2] where we choose nums1[0], nums2[1], and nums1[2].
  The sum of the numbers chosen from nums1 equals the sum of the numbers chosen from nums2: 1 + 5 = 6.
- [0, 2] where we choose nums1[0], nums1[1], and nums2[2].
  The sum of the numbers chosen from nums1 equals the sum of the numbers chosen from nums2: 1 + 2 = 3.
Note that the second and third balanced ranges are different.
In the second balanced range, we choose nums2[1] and in the third balanced range, we choose nums1[1].

Example 2:

Input: nums1 = [0,1], nums2 = [1,0]
Output: 4
Explanation: The balanced ranges are:
- [0, 0] where we choose nums1[0].
  The sum of the numbers chosen from nums1 equals the sum of the numbers chosen from nums2: 0 = 0.
- [1, 1] where we choose nums2[1].
  The sum of the numbers chosen from nums1 equals the sum of the numbers chosen from nums2: 0 = 0.
- [0, 1] where we choose nums1[0] and nums2[1].
  The sum of the numbers chosen from nums1 equals the sum of the numbers chosen from nums2: 0 = 0.
- [0, 1] where we choose nums2[0] and nums1[1].
  The sum of the numbers chosen from nums1 equals the sum of the numbers chosen from nums2: 1 = 1.

Constraints:

n == nums1.length == nums2.length
1 <= n <= 100
0 <= nums1[i], nums2[i] <= 100

Solution

Method 1 – Dynamic Programming with Prefix Sums

Intuition

The key idea is to use dynamic programming to count, for every subarray [l, r], the number of ways to pick elements from nums1 or nums2 such that the sum of picked elements from nums1 equals the sum from nums2. We use prefix sums and a DP table to efficiently compute the number of balanced ranges.

Approach

For every possible subarray [l, r]:
- Use a DP table where dp[i][d] is the number of ways to pick from the first i elements so that the difference between the sum of picked nums1 and nums2 is d.
- For each position, either pick from nums1 or nums2 and update the difference accordingly.
For each subarray, count the number of ways where the difference is zero (i.e., sums are equal).
Sum up the counts for all subarrays.
Return the answer modulo 1e9+7.

Code

C++

class Solution {
public:
    int countSubranges(vector<int>& a, vector<int>& b) {
        int n = a.size(), mod = 1e9+7, ans = 0;
        for (int l = 0; l < n; ++l) {
            unordered_map<int, int> dp;
            dp[0] = 1;
            for (int r = l; r < n; ++r) {
                unordered_map<int, int> ndp;
                for (auto& [d, cnt] : dp) {
                    ndp[d + a[r] - b[r]] = (ndp[d + a[r] - b[r]] + cnt) % mod;
                    ndp[d + b[r] - a[r]] = (ndp[d + b[r] - a[r]] + cnt) % mod;
                }
                dp = ndp;
                ans = (ans + dp[0]) % mod;
            }
        }
        return ans;
    }
};

Go

func countSubranges(a, b []int) int {
    n, mod, ans := len(a), 1_000_000_007, 0
    for l := 0; l < n; l++ {
        dp := map[int]int{0: 1}
        for r := l; r < n; r++ {
            ndp := map[int]int{}
            for d, cnt := range dp {
                ndp[d+a[r]-b[r]] = (ndp[d+a[r]-b[r]] + cnt) % mod
                ndp[d+b[r]-a[r]] = (ndp[d+b[r]-a[r]] + cnt) % mod
            }
            dp = ndp
            ans = (ans + dp[0]) % mod
        }
    }
    return ans
}

Java

class Solution {
    public int countSubranges(int[] a, int[] b) {
        int n = a.length, mod = 1_000_000_007, ans = 0;
        for (int l = 0; l < n; l++) {
            Map<Integer, Integer> dp = new HashMap<>();
            dp.put(0, 1);
            for (int r = l; r < n; r++) {
                Map<Integer, Integer> ndp = new HashMap<>();
                for (var e : dp.entrySet()) {
                    int d = e.getKey(), cnt = e.getValue();
                    ndp.put(d + a[r] - b[r], (ndp.getOrDefault(d + a[r] - b[r], 0) + cnt) % mod);
                    ndp.put(d + b[r] - a[r], (ndp.getOrDefault(d + b[r] - a[r], 0) + cnt) % mod);
                }
                dp = ndp;
                ans = (ans + dp.getOrDefault(0, 0)) % mod;
            }
        }
        return ans;
    }
}

Kotlin

class Solution {
    fun countSubranges(a: IntArray, b: IntArray): Int {
        val n = a.size
        val mod = 1_000_000_007
        var ans = 0
        for (l in 0 until n) {
            var dp = mutableMapOf(0 to 1)
            for (r in l until n) {
                val ndp = mutableMapOf<Int, Int>()
                for ((d, cnt) in dp) {
                    ndp[d + a[r] - b[r]] = (ndp.getOrDefault(d + a[r] - b[r], 0) + cnt) % mod
                    ndp[d + b[r] - a[r]] = (ndp.getOrDefault(d + b[r] - a[r], 0) + cnt) % mod
                }
                dp = ndp
                ans = (ans + dp.getOrDefault(0, 0)) % mod
            }
        }
        return ans
    }
}

Python

class Solution:
    def countSubranges(self, a: list[int], b: list[int]) -> int:
        n, mod, ans = len(a), 10**9+7, 0
        for l in range(n):
            dp = {0: 1}
            for r in range(l, n):
                ndp = {}
                for d, cnt in dp.items():
                    ndp[d + a[r] - b[r]] = (ndp.get(d + a[r] - b[r], 0) + cnt) % mod
                    ndp[d + b[r] - a[r]] = (ndp.get(d + b[r] - a[r], 0) + cnt) % mod
                dp = ndp
                ans = (ans + dp.get(0, 0)) % mod
        return ans

Rust

impl Solution {
    pub fn count_subranges(a: Vec<i32>, b: Vec<i32>) -> i32 {
        let n = a.len();
        let mut ans = 0;
        let m = 1_000_000_007;
        for l in 0..n {
            let mut dp = std::collections::HashMap::new();
            dp.insert(0, 1);
            for r in l..n {
                let mut ndp = std::collections::HashMap::new();
                for (&d, &cnt) in dp.iter() {
                    *ndp.entry(d + a[r] - b[r]).or_insert(0) = (ndp.get(&(d + a[r] - b[r])).unwrap_or(&0) + cnt) % m;
                    *ndp.entry(d + b[r] - a[r]).or_insert(0) = (ndp.get(&(d + b[r] - a[r])).unwrap_or(&0) + cnt) % m;
                }
                dp = ndp;
                ans = (ans + *dp.get(&0).unwrap_or(&0)) % m;
            }
        }
        ans
    }
}

TypeScript

class Solution {
    countSubranges(a: number[], b: number[]): number {
        const n = a.length, mod = 1e9+7;
        let ans = 0;
        for (let l = 0; l < n; l++) {
            let dp = new Map<number, number>();
            dp.set(0, 1);
            for (let r = l; r < n; r++) {
                let ndp = new Map<number, number>();
                for (let [d, cnt] of dp.entries()) {
                    ndp.set(d + a[r] - b[r], ((ndp.get(d + a[r] - b[r]) ?? 0) + cnt) % mod);
                    ndp.set(d + b[r] - a[r], ((ndp.get(d + b[r] - a[r]) ?? 0) + cnt) % mod);
                }
                dp = ndp;
                ans = (ans + (dp.get(0) ?? 0)) % mod;
            }
        }
        return ans;
    }
}

Complexity

⏰ Time complexity: O(n³·S), where S is the possible sum difference (bounded by n·max(nums1[i], nums2[i])).
🧺 Space complexity: O(S), for the DP map per subarray.