DeepSeek LeetCode 3072. 将元素分配到两个数组中 II Java实现

AI1周前发布 beixibaobao
10 0 0

这是 LeetCode 3072 题的一个 Java 实现,使用树状数组(Fenwick Tree)+ 离散化来高效维护排名和计数:

```java
class Solution {
    public int[] resultArray(int[] nums) {
        // 1. 离散化
        int[] sorted = nums.clone();
        Arrays.sort(sorted);
        int n = nums.length;
        
        // 2. 初始化两个数组
        List<Integer> arr1 = new ArrayList<>();
        List<Integer> arr2 = new ArrayList<>();
        arr1.add(nums[0]);
        arr2.add(nums[1]);
        
        // 3. 树状数组维护每个数组中元素出现的个数(按离散化后的值)
        FenwickTree tree1 = new FenwickTree(n);
        FenwickTree tree2 = new FenwickTree(n);
        
        int idx1 = getIndex(sorted, nums[0]);
        int idx2 = getIndex(sorted, nums[1]);
        tree1.add(idx1, 1);
        tree2.add(idx2, 1);
        
        // 4. 处理剩余元素
        for (int i = 2; i < n; i++) {
            int x = nums[i];
            int idx = getIndex(sorted, x);
            
            int greaterCount1 = tree1.query(n) – tree1.query(idx); // > x 的数量
            int greaterCount2 = tree2.query(n) – tree2.query(idx);
            
            if (greaterCount1 > greaterCount2) {
                arr1.add(x);
                tree1.add(idx, 1);
            } else if (greaterCount1 < greaterCount2) {
                arr2.add(x);
                tree2.add(idx, 1);
            } else {
                // 相等时,长度小的加
                if (arr1.size() <= arr2.size()) {
                    arr1.add(x);
                    tree1.add(idx, 1);
                } else {
                    arr2.add(x);
                    tree2.add(idx, 1);
                }
            }
        }
        
        // 5. 合并结果
        arr1.addAll(arr2);
        return arr1.stream().mapToInt(i -> i).toArray();
    }
    
    private int getIndex(int[] sorted, int val) {
        // 返回离散化后的索引(1-based)
        return Arrays.binarySearch(sorted, val) + 1;
    }
    
    // 树状数组模板
    class FenwickTree {
        int[] tree;
        int size;
        
        FenwickTree(int n) {
            size = n;
            tree = new int[n + 2];
        }
        
        void add(int i, int delta) {
            while (i <= size) {
                tree[i] += delta;
                i += i & -i;
            }
        }
        
        int query(int i) {
            int sum = 0;
            while (i > 0) {
                sum += tree[i];
                i -= i & -i;
            }
            return sum;
        }
    }
}
```

核心思路

1. 题意简化:维护两个数组 arr1、arr2,每来一个新元素 x,比较 arr1 中大于 x 的元素个数 与 arr2 中大于 x 的元素个数:
   · 如果前者大 → 放到 arr1
   · 如果后者大 → 放到 arr2
   · 相等 → 放到长度较小的数组
2. 高效比较:
   · 直接遍历会超时(O(n²))
   · 用树状数组快速求「大于 x 的元素个数」= 总个数 – 小于等于 x 的个数
   · 需要离散化,因为数值范围大(10⁹)
3. 离散化:
   · 排序去重(实际直接用原数组排序即可,元素可能重复)
   · binarySearch 得到排名(1‑based,便于树状数组)
4. 树状数组操作:
   · add(idx, 1):插入元素
   · query(idx):前缀和,即 ≤ idx 的元素个数
   · 大于 x 的个数 = query(n) – query(idx)

复杂度

· 时间:O(n log n)
· 空间:O(n)

© 版权声明

相关文章