将元素分配到两个数组中 II

难度: Hard

给你一个下标从 1 开始、长度为 n 的整数数组 nums 。

现定义函数 greaterCount ，使得 greaterCount(arr, val) 返回数组 arr 中 严格大于 val 的元素数量。

你需要使用 n 次操作，将 nums 的所有元素分配到两个数组 arr1 和 arr2 中。在第一次操作中，将 nums[1] 追加到 arr1 。在第二次操作中，将 nums[2] 追加到 arr2 。之后，在第 i 次操作中：

如果 greaterCount(arr1, nums[i]) > greaterCount(arr2, nums[i]) ，将 nums[i] 追加到 arr1 。
如果 greaterCount(arr1, nums[i]) < greaterCount(arr2, nums[i]) ，将 nums[i] 追加到 arr2 。
如果 greaterCount(arr1, nums[i]) == greaterCount(arr2, nums[i]) ，将 nums[i] 追加到元素数量较少的数组中。
如果仍然相等，那么将 nums[i] 追加到 arr1 。

连接数组 arr1 和 arr2 形成数组 result 。例如，如果 arr1 == [1,2,3] 且 arr2 == [4,5,6] ，那么 result = [1,2,3,4,5,6] 。

返回整数数组 result 。

示例 1：

输入：nums = [2,1,3,3]
输出：[2,3,1,3]
解释：在前两次操作后，arr1 = [2] ，arr2 = [1] 。
在第 3 次操作中，两个数组中大于 3 的元素数量都是零，并且长度相等，因此，将 nums[3] 追加到 arr1 。
在第 4 次操作中，两个数组中大于 3 的元素数量都是零，但 arr2 的长度较小，因此，将 nums[4] 追加到 arr2 。
在 4 次操作后，arr1 = [2,3] ，arr2 = [1,3] 。
因此，连接形成的数组 result 是 [2,3,1,3] 。

示例 2：

输入：nums = [5,14,3,1,2]
输出：[5,3,1,2,14]
解释：在前两次操作后，arr1 = [5] ，arr2 = [14] 。
在第 3 次操作中，两个数组中大于 3 的元素数量都是一，并且长度相等，因此，将 nums[3] 追加到 arr1 。
在第 4 次操作中，arr1 中大于 1 的元素数量大于 arr2 中的数量（2 > 1），因此，将 nums[4] 追加到 arr1 。
在第 5 次操作中，arr1 中大于 2 的元素数量大于 arr2 中的数量（2 > 1），因此，将 nums[5] 追加到 arr1 。
在 5 次操作后，arr1 = [5,3,1,2] ，arr2 = [14] 。
因此，连接形成的数组 result 是 [5,3,1,2,14] 。

示例 3：

输入：nums = [3,3,3,3]
输出：[3,3,3,3]
解释：在 4 次操作后，arr1 = [3,3] ，arr2 = [3,3] 。
因此，连接形成的数组 result 是 [3,3,3,3] 。

提示：

3 <= n <= 10⁵
1 <= nums[i] <= 10⁹

Submission

运行时间: 633 ms

内存: 41.9 MB

class Solution:
    def resultArray(self, nums: List[int]) -> List[int]:
        
        v2i = {v: i for i, v in enumerate(sorted(set(nums)), 1)}
        m = len(v2i)
        a, b = list(), list()
        t = [0] * (m + 1)
        
        for idx, x in enumerate(nums):
            i = v2i[x]
            tmp, i0 = 0, i
            while i0:
                tmp += t[i0]
                i0 &= i0 - 1
            diff = len(a) - len(b) - tmp
            if (idx == 0 or idx != 1) and (diff > 0 or diff == 0 and len(a) <= len(b)):
                a.append(x)
                while i <= m:
                    t[i] += 1
                    i += i & -i
            else:
                b.append(x)
                while i <= m:
                    t[i] -= 1
                    i += i & -i
        return a + b

Explain

该题解采用了二进制索引树（Binary Indexed Tree，BIT）来高效地维护和查询数组中特定值的计数信息。首先，创建一个值到索引的映射（v2i），将数组中的元素映射到一个连续的索引范围内，并按值排序。这个映射用于在BIT中更新和查询位置。对于数组中的每个元素，基于其在BIT中的计数值以及当前两个数组（arr1和arr2）的长度，决定将元素添加到哪一个数组。使用BIT可以快速计算出任意元素值大于当前元素值的元素数量。这种方法允许我们在对数时间内处理每个元素的插入操作，从而有效地解决问题。

时间复杂度: O(n log m)

空间复杂度: O(m + n)

class Solution:
    def resultArray(self, nums: List[int]) -> List[int]:
        # 创建值到索引的映射
        v2i = {v: i for i, v in enumerate(sorted(set(nums)), 1)}
        m = len(v2i)
        a, b = list(), list()
        t = [0] * (m + 1)
        
        for idx, x in enumerate(nums):
            i = v2i[x]
            # 查询BIT得到大于当前元素的数量
            tmp, i0 = 0, i
            while i0:
                tmp += t[i0]
                i0 &= i0 - 1
            diff = len(a) - len(b) - tmp
            # 根据条件判断将元素加入arr1或arr2
            if (idx == 0 or idx != 1) and (diff > 0 or diff == 0 and len(a) <= len(b)):
                a.append(x)
                # 更新BIT
                while i <= m:
                    t[i] += 1
                    i += i & -i
            else:
                b.append(x)
                while i <= m:
                    t[i] -= 1
                    i += i & -i
        # 返回结果数组
        return a + b

Explore

二进制索引树（Binary Indexed Tree，BIT），也被称为树状数组，主要用于高效地处理数组的前缀和查询和更新操作。在本题中，BIT 被用于维护和查询数组元素的计数信息。通过BIT，可以快速计算出任意给定元素值大于当前元素值的元素数量，这是通过在每次添加元素到数组时更新BIT，并在需要时查询所有大于当前元素的索引的计数总和来实现的。这样，每次元素插入的处理时间能够保持在对数级别（O(log n)），从而使整体解决方案更加高效。

值到索引的映射（v2i）是通过将数组中的所有不同元素进行排序和去重后，再给每一个唯一元素分配一个连续的整数索引来构建的。这种映射是必要的，因为BIT操作需要使用连续的索引范围来有效地更新和查询。如果直接使用元素的原始值作为索引，可能会导致索引过大（尤其是在元素值范围广泛时），从而增加空间复杂度。通过将元素映射到一个较小的、连续的索引范围内，我们可以使空间使用更加高效，并确保BIT的操作能够正确执行。

在这行代码中，`diff`计算的是数组`a`和`b`的长度差减去当前元素之后所有元素的数量（即`tmp`，通过BIT计算得到）。`diff`的值用于决定新元素应该添加到哪一个数组中。如果`diff`大于0，或者`diff`等于0且`a`的长度不大于`b`的长度，新元素会被添加到数组`a`中；否则，添加到数组`b`中。这种方法的目的是尽可能平衡两个数组的长度，同时考虑到元素的相对大小顺序，以达到题目的要求。