Hard

题目描述

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

定义函数 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。

数组 result 由数组 arr1 和 arr2 拼接形成。例如，如果 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]
解释：前 2 次操作后，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]

示例 3：

输入：nums = [3,3,3,3]
输出：[3,3,3,3]

提示：

  • 3 <= n <= 10^5
  • 1 <= nums[i] <= 10^9

解题思路

这道题的核心是需要高效地计算数组中严格大于某个值的元素个数，并且支持动态插入元素。

思路分析：

  1. 暴力解法：每次都遍历两个数组来计算 greaterCount，时间复杂度为 O(n²)，在 n=10⁵ 时会超时。

  2. 优化思路：需要使用支持快速查询和插入的数据结构。由于数值范围很大(10⁹)，需要先进行坐标压缩，将所有可能出现的数值映射到较小的范围内。

  3. 数据结构选择：

    • 树状数组(Binary Indexed Tree)：支持 O(log n) 的单点更新和前缀和查询
    • 线段树(Segment Tree)：同样支持 O(log n) 的区间查询和单点更新

具体实现：

  • 首先对所有数值进行离散化处理，建立数值到索引的映射
  • 使用两个树状数组分别维护 arr1 和 arr2 中各数值的出现次数
  • 对于查询大于某值的元素个数，转化为查询该值在离散化数组中位置之后的元素总数
  • 每次插入元素时，更新对应的树状数组

这种方法的时间复杂度为 O(n log n)，空间复杂度为 O(n)，可以高效解决此问题。

代码实现

class Solution {
public:
    vector<int> resultArray(vector<int>& nums) {
        int n = nums.size();
        
        // 离散化
        vector<int> sorted_nums = nums;
        sort(sorted_nums.begin(), sorted_nums.end());
        sorted_nums.erase(unique(sorted_nums.begin(), sorted_nums.end()), sorted_nums.end());
        
        int m = sorted_nums.size();
        
        auto getIndex = [&](int val) {
            return lower_bound(sorted_nums.begin(), sorted_nums.end(), val) - sorted_nums.begin() + 1;
        };
        
        // 树状数组
        vector<int> tree1(m + 1, 0), tree2(m + 1, 0);
        
        auto update = [&](vector<int>& tree, int i, int delta) {
            for (; i <= m; i += i & -i) {
                tree[i] += delta;
            }
        };
        
        auto query = [&](vector<int>& tree, int i) {
            int sum = 0;
            for (; i > 0; i -= i & -i) {
                sum += tree[i];
            }
            return sum;
        };
        
        auto greaterCount = [&](vector<int>& tree, int val) {
            int idx = getIndex(val);
            return query(tree, m) - query(tree, idx);
        };
        
        vector<int> arr1, arr2;
        
        // 第一次操作
        arr1.push_back(nums[0]);
        update(tree1, getIndex(nums[0]), 1);
        
        // 第二次操作
        arr2.push_back(nums[1]);
        update(tree2, getIndex(nums[1]), 1);
        
        // 后续操作
        for (int i = 2; i < n; i++) {
            int count1 = greaterCount(tree1, nums[i]);
            int count2 = greaterCount(tree2, nums[i]);
            
            if (count1 > count2 || (count1 == count2 && arr1.size() <= arr2.size())) {
                arr1.push_back(nums[i]);
                update(tree1, getIndex(nums[i]), 1);
            } else {
                arr2.push_back(nums[i]);
                update(tree2, getIndex(nums[i]), 1);
            }
        }
        
        // 拼接结果
        vector<int> result = arr1;
        result.insert(result.end(), arr2.begin(), arr2.end());
        
        return result;
    }
};

class Solution:
    def resultArray(self, nums: List[int]) -> List[int]:
        n = len(nums)
        
        # 离散化
        sorted_nums = sorted(list(set(nums)))
        m = len(sorted_nums)
        
        def getIndex(val):
            return bisect.bisect_left(sorted_nums, val) + 1
        
        # 树状数组
        tree1 = [0] * (m + 1)
        tree2 = [0] * (m + 1)
        
        def update(tree, i, delta):
            while i <= m:
                tree[i] += delta
                i += i & -i
        
        def query(tree, i):
            s = 0
            while i > 0:
                s += tree[i]
                i -= i & -i
            return s
        
        def greaterCount(tree, val):
            idx = getIndex(val)
            return query(tree, m) - query(tree, idx)
        
        arr1, arr2 = [], []
        
        # 第一次操作
        arr1.append(nums[0])
        update(tree1, getIndex(nums[0]), 1)
        
        # 第二次操作
        arr2.append(nums[1])
        update(tree2, getIndex(nums[1]), 1)
        
        # 后续操作
        for i in range(2, n):
            count1 = greaterCount(tree1, nums[i])
            count2 = greaterCount(tree2, nums[i])
            
            if count1 > count2 or (count1 == count2 and len(arr1) <= len(arr2)):
                arr1.append(nums[i])
                update(tree1, getIndex(nums[i]), 1)
            else:
                arr2.append(nums[i])
                update(tree2, getIndex(nums[i]), 1)
        
        return arr1 + arr2

public class Solution {
    public int[] ResultArray(int[] nums) {
        int n = nums.Length;
        
        // 离散化
        var sortedNums = nums.Distinct().OrderBy(x => x).ToArray();
        int m = sortedNums.Length;
        var indexMap = sortedNums.Select((val, idx) => new { val, idx }).ToDictionary(x => x.val, x => x.idx + 1);
        
        int GetIndex(int val) => indexMap[val];
        
        // 树状数组
        int[] tree1 = new int[m + 1];
        int[] tree2 = new int[m + 1];
        
        void Update(int[] tree, int i, int delta) {
            for (; i <= m; i += i & -i) {
                tree[i] += delta;
            }
        }
        
        int Query(int[] tree, int i) {
            int sum = 0;
            for (; i > 0; i -= i & -i) {
                sum += tree[i];
            }
            return sum;
        }
        
        int GreaterCount(int[] tree, int val) {
            int idx = GetIndex(val);
            return Query(tree, m) - Query(tree, idx);
        }
        
        var arr1 = new List<int>();
        var arr2 = new List<int>();
        
        // 第一次操作
        arr1.Add(nums[0]);
        Update(tree1, GetIndex(nums[0]), 1);
        
        // 第二次操作
        arr2.Add(nums[1]);
        Update(tree2, GetIndex(nums[1]), 1);
        
        // 后续操作
        for (int i = 2; i < n; i++) {
            int count1 = GreaterCount(tree1, nums[i]);
            int count2 = GreaterCount(tree2, nums[i]);
            
            if (count1 > count2 || (count1 == count2 && arr1.Count <= arr2.Count)) {
                arr1.Add(nums[i]);
                Update(tree1, GetIndex(nums[i]), 1);
            } else {
                arr2.Add(nums[i]);
                Update(tree2, GetIndex(nums[i]), 1);
            }
        }
        
        var result = new int[n];
        arr1.CopyTo(result, 0);
        arr2.CopyTo(result, arr1.Count);
        
        return result;
    }
}

var resultArray = function(nums) {
    class BIT {
        constructor(size) {
            this.size = size;
            this.tree = new Array(size + 1).fill(0);
        }
        
        update(idx, val) {
            for (let i = idx; i <= this.size; i += i & (-i)) {
                this.tree[i] += val;
            }
        }
        
        query(idx) {
            let sum = 0;
            for (let i = idx; i > 0; i -= i & (-i)) {
                sum += this.tree[i];
            }
            return sum;
        }
        
        rangeQuery(left, right) {
            return this.query(right) - this.query(left - 1);
        }
    }
    
    let sorted = [...new Set(nums)].sort((a, b) => a - b);
    let coordMap = new Map();
    sorted.forEach((val, idx) => coordMap.set(val, idx + 1));
    
    let arr1 = [nums[0]];
    let arr2 = [nums[1]];
    let bit1 = new BIT(sorted.length);
    let bit2 = new BIT(sorted.length);
    
    bit1.update(coordMap.get(nums[0]), 1);
    bit2.update(coordMap.get(nums[1]), 1);
    
    for (let i = 2; i < nums.length; i++) {
        let coord = coordMap.get(nums[i]);
        let count1 = bit1.rangeQuery(coord + 1, sorted.length);
        let count2 = bit2.rangeQuery(coord + 1, sorted.length);
        
        if (count1 > count2 || (count1 === count2 && arr1.length <= arr2.length)) {
            arr1.push(nums[i]);
            bit1.update(coord, 1);
        } else {
            arr2.push(nums[i]);
            bit2.update(coord, 1);
        }
    }
    
    return arr1.concat(arr2);
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型复杂度说明
时间复杂度O(n log n)离散化排序 O(n log n) + n次操作每次 O(log n)
空间复杂度O(n)树状数组和结果数组的空间

相关题目