[每日一题] 3072.将元素分配到两个数组中II

2024 年 06 月 05 日

90 次浏览

1600字数

> https://leetcode.cn/problems/distribute-elements-into-two-arrays-ii/description/?envType=daily-question&envId=2024-05-06

要实现 `greaterCount` 函数，就需要快速找到严格大于 `val` 的元素数量，并且支持快速插入元素，因此可以使用树状数组来解决。

我们只关注数组元素的大小关系，因此可以将数组离散化处理。然后初始化两个结果数组和对应的树，依次遍历 `nums` 数组中的数并将其离散化后的索引加入树中，最后拼接即可。

```c++
class Solution {
public:
    std::vector<int> resultArray(std::vector<int>& nums) {
        class BinaryIndexTree {
        private:
            std::vector<int> tree;

public:
            BinaryIndexTree(int n) : tree(n + 1) {}

void add(int x) {
                while (x < tree.size()) {
                    tree[x]++;
                    x += x & -x;
                }
            }

int getSum(int x) {
                int sum = 0;
                while (x > 0) {
                    sum += tree[x];
                    x -= x & -x;
                }
                return sum;
            }
        };

const int n = nums.size();

std::vector<int> sorted(nums);
        std::sort(sorted.begin(), sorted.end());

std::unordered_map<int, int> index;
        for (int i = 0; i < n; i++)
            index[sorted[i]] = i + 1;

std::vector<int> arr1 = {nums[0]}, arr2 = {nums[1]};
        BinaryIndexTree tree1(n), tree2(n);
        tree1.add(index[nums[0]]);
        tree2.add(index[nums[1]]);

for (int i = 2; i < n; i++) {
            int cnt1 = arr1.size() - tree1.getSum(index[nums[i]]);
            int cnt2 = arr2.size() - tree2.getSum(index[nums[i]]);

if (cnt1 > cnt2 || (cnt1 == cnt2 && arr1.size() <= arr2.size())) {
                arr1.push_back(nums[i]);
                tree1.add(index[nums[i]]);
            } else {
                arr2.push_back(nums[i]);
                tree2.add(index[nums[i]]);
            }
        }

arr1.insert(arr1.end(), arr2.begin(), arr2.end());
        return arr1;
    }
};
```

- 时间复杂度：$O(nlogn)$ ，其中 $n$ 为 `nums` 数组的长度。
- 空间复杂度：$O(n)$ 。

最后修改：2024 年 06 月 05 日