2143. 在两个数组的区间中选取数字

English Version

题目描述

给你两个 下标从 0 开始，长度为 n 的整数数组 nums1 和 nums2。

如果一个区间 [l, r] （包含左右端点，0 <= l <= r < n）满足下列条件，那么这个区间就是 平衡 的：

• 对每个在区间 [l, r] 范围内的 i，你需要选取 nums1[i] 或者 nums2[i]；
• 从 nums1 中选取的数字和与从 nums2 中选取的数字和相等。（如果你没有从某个数组中选取任何数字，那么数字和被视为 0）。

如果两个 平衡 的区间 [l1, r1] 和 [l2, r2] 满足下列条件之一，那么它们就是 不同 的：

• l1 != l2
• r1 != r2
• 两个区间中的数字选取情况不同（也就是说，存在至少一个 i，使得在第一个区间中，nums1[i] 被选中, 而在第二个区间中，nums2[i] 被选中，或者相反的情况）。

请返回 不同 的平衡的区间数目。由于答案可能很大，请返回答案 模 109+7 的结果。

 

示例 1:

输入: nums1 = [1,2,5], nums2 = [2,6,3]
输出: 3
解释: 平衡的区间有:
- [0, 1], 我们选取 nums2[0] 和 nums2[1]。
  从 nums1 中选取的数字和与从 nums2 中选取的数字和相等: 2 = 2.
- [0, 2], 我们选取 nums1[0], nums2[1] 和 nums1[2]。
  从 nums1 中选取的数字和与从 nums2 中选取的数字和相等: 1 + 5 = 6。
- [0, 2], 我们选取 nums1[0], nums1[1] 和 nums2[2]。
  从 nums1 中选取的数字和与从 nums2 中选取的数字和相等: 1 + 2 = 3。
注意第二个区间和第三个区间是不同的。
因为在第二个平衡的区间中，我们选取了 nums2[1]，但是在第三个平衡的区间中，我们选取了 nums1[1]。

示例 2：

输入: nums1 = [0,1], nums2 = [1,0]
输出: 4
解释: 平衡的区间有:
- [0, 0], 我们选取 nums1[0]。
  从 nums1 中选取的数字和与从 nums2 中选取的数字和相等: 0 = 0。
- [1, 1], 我们选取 nums2[1]。
  从 nums1 中选取的数字和与从 nums2 中选取的数字和相等: 0 = 0。
- [0, 1], 我们选取 nums1[0] 和 nums2[1]。
  从 nums1 中选取的数字和与从 nums2 中选取的数字和相等: 0 = 0。
- [0, 1], 我们选取 nums2[0] 和 nums1[1]。
  从 nums1 中选取的数字和与从 nums2 中选取的数字和相等: 1 = 1。

 

提示:

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

解法

方法一：动态规划

我们定义 $f[i][j]$ 表示以第 $i$ 个元素结尾，且从 $nums1$ 中选取的数字和与从 $nums2$ 中选取的数字和之差为 $j$ 的平衡区间的个数。由于差值可能为负数，因此，我们统一将 $j$ 加上 $s_2 = \sum_{k=0}^{n-1}nums2[k]$，这样就可以保证 $j$ 为非负整数。

考虑 $f[i][j]$，我们可以单独将第 $i$ 个元素视为一个区间，那么 $f[i][nums1[i] + s_2]$ 和 $f[i][-nums2[i] + s_2]$ 都会增加 $1$。此外，如果 $i \gt 0$，我们也可以将第 $i$ 个元素添加到前面的某个区间中，我们在 $[0, s_1 + s_2]$ 范围内枚举 $j$，如果 $j \geq a$，那么 $f[i][j]$ 会增加 $f[i - 1][j - a]$，如果 $j + b \leq s_1 + s_2$，那么 $f[i][j]$ 会增加 $f[i - 1][j + b]$。

答案为 $\sum_{i=0}^{n-1}f[i][s_2]$。

时间复杂度 $O(n \times M)$，空间复杂度 $O(n \times M)$。其中 $n$ 和 $M$ 分别为数组 $nums1$ 的长度以及数字和的最大值。

Python3

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

Java

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;
    }
}

C++

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;
    }
};

Go

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
}

TypeScript

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;
}

...