DeepSeek LeetCode 3072. 将元素分配到两个数组中 II Java实现
这是 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)