树状数组+离散化，LeetCode 3072. 将元素分配到两个数组中 II

一、题目

1、题目描述

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

2、原题链接

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

---

二、解题报告

1、思路分析

我们只需要能够实现快速查询arr1和arr2内大于给定数字的数字个数即可

考虑线段树、树状数组、莫队......

这里使用两个树状数组维护

Python可以直接使用有序列表SortedList

关于树状数组：前缀和的动态维护------树状数组[C/C++]_树状数组维护前缀和-CSDN博客

2、复杂度

时间复杂度： O(NlogN) 空间复杂度：O(N)

3、代码详解

​cpp

template<class T = int>
class Fenwick {
private:
    std::vector<T> tr;
    T n;
public:
    Fenwick(T _n) : tr(_n + 1), n(_n) {}
    void add(T x, T k) {
        for (; x <= n; x += x & -x) 
            tr[x] += k;
    }
    T query(int x) const {
        T res = 0;
        for (; x; x &= x - 1)
            res += tr[x];
        return res;
    }
};
class Solution {
public:
    vector<int> resultArray(vector<int>& nums) {
        std::vector<int> idx(nums), arr1 { nums[0] }, arr2 { nums[1] };
        std::sort(idx.begin(), idx.end());
        idx.erase(std::unique(idx.begin(), idx.end()), idx.end());
        int n = nums.size(), m = idx.size();
        Fenwick<int> st1(m), st2(m);
        st1.add(std::lower_bound(idx.begin(), idx.end(), nums[0]) - idx.begin() + 1, 1);
        st2.add(std::lower_bound(idx.begin(), idx.end(), nums[1]) - idx.begin() + 1, 1);
        for (int i = 2; i < n; i ++ ) {
            int x = std::lower_bound(idx.begin(), idx.end(), nums[i]) - idx.begin() + 1;
            int cl = st1.query(m) - st1.query(x);
            int cr = st2.query(m) - st2.query(x);
            if (cl < cr) {
                st2.add(x, 1);
                arr2.push_back(nums[i]);
            }
            else if (cl > cr) {
                st1.add(x, 1);
                arr1.push_back(nums[i]);
            }
            else if (cl == cr) {
                if (arr1.size() <= arr2.size()) {
                    st1.add(x, 1);
                    arr1.push_back(nums[i]);
                }
                else {
                    st2.add(x, 1);
                    arr2.push_back(nums[i]);
                }
            }
        }
        arr1.insert(arr1.end(), arr2.begin(), arr2.end());
        return arr1;
    }
};

py3

from sortedcontainers import SortedList
class Solution:
    def resultArray(self, nums: List[int]) -> List[int]:
        arr1, arr2 = [nums[0]], [nums[1]]
        sl, sr = SortedList(arr1), SortedList(arr2)
        
        for x in nums[2::]:
            cl = len(sl) - sl.bisect_right(x)
            cr = len(sr) - sr.bisect_right(x)
            
            if cl < cr:
                sr.add(x)
                arr2.append(x)
            elif cl > cr:
                sl.add(x)
                arr1.append(x)
            elif cl == cr:
                if len(sl) <= len(sr):
                    sl.add(x)
                    arr1.append(x)
                else:
                    sr.add(x)
                    arr2.append(x)
        return arr1 + arr2