3072. 将元素分配到两个数组中 II Rust 线段树 + 离散化

题目

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