2763-所有子数组中不平衡数字之和

Raphael Liu Lv10
  2023-09-13 00:47:26 2023-09-13 00:47:26 Created   2023-09-13 16:06:26 2023-09-13 16:06:26 Updated    

一个长度为 n 下标从 0 开始的整数数组 arr 的 不平衡数字 定义为，在 sarr = sorted(arr)
数组中，满足以下条件的下标数目：

  • 0 <= i < n - 1 ，和
  • sarr[i+1] - sarr[i] > 1

这里，sorted(arr) 表示将数组 arr 排序后得到的数组。

给你一个下标从 0 开始的整数数组 nums ，请你返回它所有 子数组 的 不平衡数字 之和。

子数组指的是一个数组中连续一段 非空 的元素序列。

示例 1：

**输入：** nums = [2,3,1,4]
**输出：** 3
**解释：** 总共有 3 个子数组有非 0 不平衡数字：
- 子数组 [3, 1] ，不平衡数字为 1 。
- 子数组 [3, 1, 4] ，不平衡数字为 1 。
- 子数组 [1, 4] ，不平衡数字为 1 。
其他所有子数组的不平衡数字都是 0 ，所以所有子数组的不平衡数字之和为 3 。

示例 2：

**输入：** nums = [1,3,3,3,5]
**输出：** 8
**解释：** 总共有 7 个子数组有非 0 不平衡数字：
- 子数组 [1, 3] ，不平衡数字为 1 。
- 子数组 [1, 3, 3] ，不平衡数字为 1 。
- 子数组 [1, 3, 3, 3] ，不平衡数字为 1 。
- 子数组 [1, 3, 3, 3, 5] ，不平衡数字为 2 。
- 子数组 [3, 3, 3, 5] ，不平衡数字为 1 。
- 子数组 [3, 3, 5] ，不平衡数字为 1 。
- 子数组 [3, 5] ，不平衡数字为 1 。
其他所有子数组的不平衡数字都是 0 ，所以所有子数组的不平衡数字之和为 8 。

提示：

  • 1 <= nums.length <= 1000
  • 1 <= nums[i] <= nums.length

方法一：枚举

由于 n 至多为 1000，我们可以从左到右枚举子数组左端点 i，然后从 i+1 开始向右枚举子数组右端点 j。一边枚举 j，一边维护不平衡度 cnt：

  • 如果 x=\textit{nums}[j] 之前出现过，那么子数组排序后必然会和另一个 x 相邻，cnt 不变；
  • 如果 x=\textit{nums}[j] 之前没出现过，那么看 x-1 和 x+1 是否出现过：
    • 都没有，cnt 加一；
    • 只有一个，cnt 不变；
    • 两个都有，cnt 减一。

遍历过程中，累加 cnt，即为答案。

[sol-Python3]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution:
    def sumImbalanceNumbers(self, nums: List[int]) -> int:
        ans, n = 0, len(nums)
        for i, x in enumerate(nums):
            vis = [False] * (n + 2)
            vis[x] = True
            cnt = 0
            for j in range(i + 1, n):
                x = nums[j]
                if not vis[x]:
                    cnt += 1 - vis[x - 1] - vis[x + 1]
                    vis[x] = True
                ans += cnt
        return ans
[sol-Java]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
    public int sumImbalanceNumbers(int[] nums) {
        int ans = 0, n = nums.length;
        var vis = new boolean[n + 2];
        for (int i = 0; i < n; i++) {
            Arrays.fill(vis, false);
            vis[nums[i]] = true;
            int cnt = 0;
            for (int j = i + 1; j < n; j++) {
                int x = nums[j];
                if (!vis[x]) {
                    cnt++;
                    if (vis[x - 1]) cnt--;
                    if (vis[x + 1]) cnt--;
                    vis[x] = true;
                }
                ans += cnt;
            }
        }
        return ans;
    }
}
[sol-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
public:
    int sumImbalanceNumbers(vector<int> &nums) {
        int ans = 0, n = nums.size();
        bool vis[n + 2];
        for (int i = 0; i < n; i++) {
            memset(vis, 0, sizeof(vis));
            vis[nums[i]] = true;
            int cnt = 0;
            for (int j = i + 1; j < n; j++) {
                int x = nums[j];
                if (!vis[x]) {
                    cnt += 1 - vis[x - 1] - vis[x + 1];
                    vis[x] = true;
                }
                ans += cnt;
            }
        }
        return ans;
    }
};
[sol-Go]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
func sumImbalanceNumbers(nums []int) (ans int) {
	n := len(nums)
	for i, x := range nums {
		vis := make([]int, n+2)
		vis[x] = 1
		cnt := 0
		for j := i + 1; j < n; j++ {
			if x := nums[j]; vis[x] == 0 {
				cnt += 1 - vis[x-1] - vis[x+1]
				vis[x] = 1
			}
			ans += cnt
		}
	}
	return
}

复杂度分析

  • 时间复杂度：\mathcal{O}(n^2)，其中 n 为 nums 的长度。
  • 空间复杂度：\mathcal{O}(n)。

方法二：贡献法

视频讲解见【周赛 352】 第四题。

[sol-Python3]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def sumImbalanceNumbers(self, nums: List[int]) -> int:
        n = len(nums)
        right = [0] * n  # nums[i] 右侧的 x 和 x-1 的最近下标（不存在时为 n）
        idx = [n] * (n + 1)
        for i in range(n - 1, -1, -1):
            x = nums[i]
            right[i] = min(idx[x], idx[x - 1])
            idx[x] = i

        ans = 0
        idx = [-1] * (n + 1)
        for i, (x, r) in enumerate(zip(nums, right)):
            # 统计 x 能产生多少贡献
            ans += (i - idx[x - 1]) * (r - i)  # 子数组左端点个数 * 子数组右端点个数
            idx[x] = i
        # 上面计算的时候，每个子数组的最小值必然可以作为贡献，而这是不合法的
        # 所以每个子数组都多算了 1 个不合法的贡献
        return ans - n * (n + 1) // 2
[sol-Java]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution {
    public int sumImbalanceNumbers(int[] nums) {
        int n = nums.length;
        var right = new int[n];
        var idx = new int[n + 1];
        Arrays.fill(idx, n);
        for (int i = n - 1; i >= 0; i--) {
            int x = nums[i];
            // right[i] 表示 nums[i] 右侧的 x 和 x-1 的最近下标（不存在时为 n）
            right[i] = Math.min(idx[x], idx[x - 1]);
            idx[x] = i;
        }

        int ans = 0;
        Arrays.fill(idx, -1);
        for (int i = 0; i < n; i++) {
            int x = nums[i];
            // 统计 x 能产生多少贡献
            ans += (i - idx[x - 1]) * (right[i] - i); // 子数组左端点个数 * 子数组右端点个数
            idx[x] = i;
        }
        // 上面计算的时候，每个子数组的最小值必然可以作为贡献，而这是不合法的
        // 所以每个子数组都多算了 1 个不合法的贡献
        return ans - n * (n + 1) / 2;
    }
}
[sol-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
public:
    int sumImbalanceNumbers(vector<int> &nums) {
        int n = nums.size(), right[n], idx[n + 1];
        fill(idx, idx + n + 1, n);
        for (int i = n - 1; i >= 0; i--) {
            int x = nums[i];
            // right[i] 表示 nums[i] 右侧的 x 和 x-1 的最近下标（不存在时为 n）
            right[i] = min(idx[x], idx[x - 1]);
            idx[x] = i;
        }

        int ans = 0;
        memset(idx, -1, sizeof(idx));
        for (int i = 0; i < n; i++) {
            int x = nums[i];
            // 统计 x 能产生多少贡献
            ans += (i - idx[x - 1]) * (right[i] - i); // 子数组左端点个数 * 子数组右端点个数
            idx[x] = i;
        }
        // 上面计算的时候，每个子数组的最小值必然可以作为贡献，而这是不合法的
        // 所以每个子数组都多算了 1 个不合法的贡献
        return ans - n * (n + 1) / 2;
    }
};
[sol-Go]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
func sumImbalanceNumbers(nums []int) (ans int) {
	n := len(nums)
	right := make([]int, n)
	idx := make([]int, n+1)
	for i := range idx {
		idx[i] = n
	}
	for i := n - 1; i >= 0; i-- {
		x := nums[i]
		// right[i] 表示 nums[i] 右侧的 x 和 x-1 的最近下标（不存在时为 n）
		right[i] = min(idx[x], idx[x-1])
		idx[x] = i
	}

	for i := range idx {
		idx[i] = -1
	}
	for i, x := range nums {
		// 统计 x 能产生多少贡献
		ans += (i - idx[x-1]) * (right[i] - i) // 子数组左端点个数 * 子数组右端点个数
		idx[x] = i
	}
	// 上面计算的时候，每个子数组的最小值必然可以作为贡献，而这是不合法的
	// 所以每个子数组都多算了 1 个不合法的贡献
	return ans - n*(n+1)/2
}

func min(a, b int) int { if b < a { return b }; return a }

复杂度分析

  • 时间复杂度：\mathcal{O}(n)，其中 n 为 nums 的长度。
  • 空间复杂度：\mathcal{O}(n)。

思考题

把 sarr[i+1] - sarr[i] > 1 改成 sarr[i+1] - sarr[i] > k 要怎么做？欢迎在评论区发表你的思路。

 Comments
On this page
2763-所有子数组中不平衡数字之和
  1. 方法一：枚举
  2. 方法二：贡献法
  3. 思考题