题目
给你一个下标从 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
。
示例
示例 1
输入:nums = [2,1,3,3]
输出:[2,3,1,3]
解释:在前两次操作后,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]
解释:在前两次操作后,arr1 = [5] ,arr2 = [14] 。
在第 3 次操作中,两个数组中大于 3 的元素数量都是一,并且长度相等,因此,将 nums[3] 追加到 arr1 。
在第 4 次操作中,arr1 中大于 1 的元素数量大于 arr2 中的数量(2 > 1),因此,将 nums[4] 追加到 arr1 。
在第 5 次操作中,arr1 中大于 2 的元素数量大于 arr2 中的数量(2 > 1),因此,将 nums[5] 追加到 arr1 。
在 5 次操作后,arr1 = [5,3,1,2] ,arr2 = [14] 。
因此,连接形成的数组 result 是 [5,3,1,2,14] 。
示例 3
输入:nums = [3,3,3,3]
输出:[3,3,3,3]
解释:在 4 次操作后,arr1 = [3,3] ,arr2 = [3,3] 。
因此,连接形成的数组 result 是 [3,3,3,3] 。
提示:
3 <= n <= 105
1 <= nums[i] <= 109
思路
离散化 + 线段树,由于基础的线段树可以AC,不再使用懒标记去优化。
AC代码
use std::collections::HashMap;
struct Tree{
tree: Vec<i32>
}
impl Tree {
pub fn new(len: usize) -> Self {
Tree {
tree: vec![0; len]
}
}
/**
* 更新
**/
pub fn update(&mut self, node: usize, sign_idx: usize, l: usize, r: usize) {
if l == r {
self.tree[node] += 1;
return;
}
let mid: usize = l + r >> 1;
let l_node: usize = (node << 1) + 1;
let r_node: usize = l_node + 1;
if sign_idx <= mid {
self.update(l_node, sign_idx, l, mid);
} else {
self.update(r_node, sign_idx, mid + 1, r);
}
self.tree[node] = self.tree[l_node] + self.tree[r_node];
}
/**
* 查询
**/
pub fn query(&mut self, node: usize, l: usize, r: usize, start: usize, end: usize) -> i32 {
if l > end || r < start {
return 0;
}
if l >= start && r <= end {
return self.tree[node];
}
let mid: usize = l + r >> 1;
let l_node: usize = (node << 1) + 1;
let r_node: usize = l_node + 1;
self.query(l_node, l, mid, start, end) + self.query(r_node, mid + 1, r, start, end)
}
}
impl Solution {
pub fn result_array(v: Vec<i32>) -> Vec<i32> {
let len: usize = v.len();
let tree_len: usize = len << 2;
let mut tree1: Tree = Tree::new(tree_len);
let mut tree2: Tree = Tree::new(tree_len);
let mut arr1: Vec<i32> = vec![v[0]];
let mut arr2: Vec<i32> = vec![v[1]];
let mut mp: HashMap<usize, usize> = HashMap::new();
let mut cp_v: Vec<i32> = v.clone();
cp_v.sort();
for (idx, tem) in cp_v.iter().enumerate() {
mp.insert(*tem as usize, idx);
}
if let Some(tem_val) = mp.get(&(v[0] as usize)) {
tree1.update(0, *tem_val, 0, len - 1);
}
if let Some(tem_val) = mp.get(&(v[1] as usize)) {
tree2.update(0, *tem_val, 0, len - 1);
}
for idx in 2 .. len {
let val: i32 = v[idx];
let mut mp_val: usize = 0;
if let Some(tem_val) = mp.get(&(val as usize)) {
mp_val = *tem_val;
}
let s1: i32 = tree1.query(0, 0, len - 1, mp_val + 1, len - 1);
let s2: i32 = tree2.query(0, 0, len - 1, mp_val + 1, len - 1);
if s1 > s2 {
arr1.push(val);
tree1.update(0, mp_val, 0, len - 1);
continue;
}
if s2 > s1 {
arr2.push(val);
tree2.update(0, mp_val, 0, len - 1);
continue;
}
let len1: usize = arr1.len();
let len2: usize = arr2.len();
if len1 <= len2 {
arr1.push(val);
tree1.update(0, mp_val, 0, len - 1);
} else {
arr2.push(val);
tree2.update(0, mp_val, 0, len - 1);
}
}
arr1.extend(arr2);
arr1
}
}