Question

2143. Choose Numbers From Two Arrays in Range

中文文档

Description

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 of different ranges that are balanced. _Since the answer may be very large, return it modulo 109 + 7_._

 

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

Solutions

Solution 1

class Solution:
  def countSubranges(self, nums1: List[int], nums2: List[int]) -> int:
    n = len(nums1)
    s1, s2 = sum(nums1), sum(nums2)
    f = [[0] * (s1 + s2 + 1) for _ in range(n)]
    ans = 0
    mod = 10**9 + 7
    for i, (a, b) in enumerate(zip(nums1, nums2)):
      f[i][a + s2] += 1
      f[i][-b + s2] += 1
      if i:
        for j in range(s1 + s2 + 1):
          if j >= a:
            f[i][j] = (f[i][j] + f[i - 1][j - a]) % mod
          if j + b < s1 + s2 + 1:
            f[i][j] = (f[i][j] + f[i - 1][j + b]) % mod
      ans = (ans + f[i][s2]) % mod
    return ans

class Solution {
  public int countSubranges(int[] nums1, int[] nums2) {
    int n = nums1.length;
    int s1 = Arrays.stream(nums1).sum();
    int s2 = Arrays.stream(nums2).sum();
    int[][] f = new int[n][s1 + s2 + 1];
    int ans = 0;
    final int mod = (int) 1e9 + 7;
    for (int i = 0; i < n; ++i) {
      int a = nums1[i], b = nums2[i];
      f[i][a + s2]++;
      f[i][-b + s2]++;
      if (i > 0) {
        for (int j = 0; j <= s1 + s2; ++j) {
          if (j >= a) {
            f[i][j] = (f[i][j] + f[i - 1][j - a]) % mod;
          }
          if (j + b <= s1 + s2) {
            f[i][j] = (f[i][j] + f[i - 1][j + b]) % mod;
          }
        }
      }
      ans = (ans + f[i][s2]) % mod;
    }
    return ans;
  }
}

class Solution {
public:
  int countSubranges(vector<int>& nums1, vector<int>& nums2) {
    int n = nums1.size();
    int s1 = accumulate(nums1.begin(), nums1.end(), 0);
    int s2 = accumulate(nums2.begin(), nums2.end(), 0);
    int f[n][s1 + s2 + 1];
    memset(f, 0, sizeof(f));
    int ans = 0;
    const int mod = 1e9 + 7;
    for (int i = 0; i < n; ++i) {
      int a = nums1[i], b = nums2[i];
      f[i][a + s2]++;
      f[i][-b + s2]++;
      if (i) {
        for (int j = 0; j <= s1 + s2; ++j) {
          if (j >= a) {
            f[i][j] = (f[i][j] + f[i - 1][j - a]) % mod;
          }
          if (j + b <= s1 + s2) {
            f[i][j] = (f[i][j] + f[i - 1][j + b]) % mod;
          }
        }
      }
      ans = (ans + f[i][s2]) % mod;
    }
    return ans;
  }
};

func countSubranges(nums1 []int, nums2 []int) (ans int) {
  n := len(nums1)
  s1, s2 := sum(nums1), sum(nums2)
  f := make([][]int, n)
  for i := range f {
    f[i] = make([]int, s1+s2+1)
  }
  const mod int = 1e9 + 7
  for i, a := range nums1 {
    b := nums2[i]
    f[i][a+s2]++
    f[i][-b+s2]++
    if i > 0 {
      for j := 0; j <= s1+s2; j++ {
        if j >= a {
          f[i][j] = (f[i][j] + f[i-1][j-a]) % mod
        }
        if j+b <= s1+s2 {
          f[i][j] = (f[i][j] + f[i-1][j+b]) % mod
        }
      }
    }
    ans = (ans + f[i][s2]) % mod
  }
  return
}

func sum(nums []int) (ans int) {
  for _, x := range nums {
    ans += x
  }
  return
}

function countSubranges(nums1: number[], nums2: number[]): number {
  const n = nums1.length;
  const s1 = nums1.reduce((a, b) => a + b, 0);
  const s2 = nums2.reduce((a, b) => a + b, 0);
  const f: number[][] = Array(n)
    .fill(0)
    .map(() => Array(s1 + s2 + 1).fill(0));
  const mod = 1e9 + 7;
  let ans = 0;
  for (let i = 0; i < n; ++i) {
    const [a, b] = [nums1[i], nums2[i]];
    f[i][a + s2]++;
    f[i][-b + s2]++;
    if (i) {
      for (let j = 0; j <= s1 + s2; ++j) {
        if (j >= a) {
          f[i][j] = (f[i][j] + f[i - 1][j - a]) % mod;
        }
        if (j + b <= s1 + s2) {
          f[i][j] = (f[i][j] + f[i - 1][j + b]) % mod;
        }
      }
    }
    ans = (ans + f[i][s2]) % mod;
  }
  return ans;
}