leetcode 每日1题 #3072 将元素分配到两个数组中 II

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

题目

给你一个下标从 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 。

题解

数据结构

树状数组，也叫二叉索引树BinaryIndexedTree，Fenwick树，是一种用于处理前缀和查询和更新的高效数据结构。它的主要优势在于它能够在对数时间内进行更新和查询操作。

class BinaryIndexedTree {
private:
    vector<int> tree;

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

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

    int get(int i) {
        int res = 0;
        while (i > 0) {
            res += tree[i];
            i &= i - 1;
        }
        return res;
    }
};

工作原理

1. 更新操作（add）：
   • 从给定的索引开始，逐步更新数组tree中的累加值。
   • 使用i += i & -i来找到下一个需要更新的位置。这个操作利用了二进制表示的特性，使得更新过程能够在对数时间内完成。从数值上看，该操作是i加上i二进制格式最后一位1所代表的数，比如6，二进制 0b110，加上0b10也就是2，得到8(0b1000)
2. 查询操作（get）：
   • 从给定的索引开始，逐步累加数组tree中的值。
   • 使用i &= i - 1来找到上一个需要累加的位置。这也是利用了二进制表示的特性，使得查询过程能够在对数时间内完成。从数值上看，该操作是把i二进制格式最后一位1置0，比如7，二进制 0b111，置0后是0b110，也就是6。
     

如上图，如果A[5]的值增加了1，5+(5&(-5))=6，6+(6&(-6))=8，C[5]、C[6]、C[8]的值都加1。写成二进制格式分别是C[0b0101]、C[0b0110]、C[0b1000]。
如果想得到get(7)的值，7&(7-1)=6，6&(6-1)=4，get(7)=C[7]+C[6]+C[4]。写成二进制格式是get(7)=C[0b111]+C[0b110]+C[0b100]。

算法

以 case nums = [5,14,3,1,2]为例
先排序，排序后，计算每个数的位次

        int n = nums.size();
        vector<int> sortedNums = nums;
        sort(sortedNums.begin(), sortedNums.end());
        unordered_map<int, int> index;
        for (int i = 0; i < n; ++i) {
            index[sortedNums[i]] = i + 1;
        }

排序后结果

初始化数组

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

tree1数组add元素5，index值4

tree2数组add元素14，index值5

主循环

        for (int i = 2; i < n; ++i) {
            int count1 = arr1.size() - tree1.get(index[nums[i]]);
            int count2 = arr2.size() - tree2.get(index[nums[i]]);
            if (count1 > count2 || (count1 == count2 && 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]]);
            }
        }

元素3的index值为3
5>3，tree1.get(3)返回c[3]+c[2]=0
41>3，tree2.get(3)返回c[3]+c[2]=0
得到count1=count2=0
插入arr1，tree1.add(3)

元素1，2 同理，插入arr1。


最后我们得到arr1 = [5, 3, 1, 2]，arr2 = [14]

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

拼接即为所求。

源码

class BinaryIndexedTree {
private:
    vector<int> tree;

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

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

    int get(int i) {
        int res = 0;
        while (i > 0) {
            res += tree[i];
            i &= i - 1;
        }
        return res;
    }
};

class Solution {
public:
    vector<int> resultArray(vector<int>& nums) {
        int n = nums.size();
        vector<int> sortedNums = nums;
        sort(sortedNums.begin(), sortedNums.end());
        unordered_map<int, int> index;
        for (int i = 0; i < n; ++i) {
            index[sortedNums[i]] = i + 1;
        }

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

        for (int i = 2; i < n; ++i) {
            int count1 = arr1.size() - tree1.get(index[nums[i]]);
            int count2 = arr2.size() - tree2.get(index[nums[i]]);
            if (count1 > count2 || (count1 == count2 && 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;
    }
};

作者：力扣官方题解
链接：https://leetcode.cn/problems/distribute-elements-into-two-arrays-ii/solutions/2796537/jiang-yuan-su-fen-pei-dao-liang-ge-shu-z-d5mh/
来源：力扣（LeetCode）
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

复杂度

• 时间复杂度：O(nlogn)，其中 n 表示 nums 数组的长度。
• 空间复杂度：O(n)。