Luobuyu's Blog

0%

100246. 将元素分配到两个数组中 II

发表于 分类于 ， ， 热度： ℃ 评论：
本文字数： 6.7k 阅读时长 ≈ 6 分钟
使用离散化加线段树，或者使用python的SortedList。

100246.将元素分配到两个数组中 II

题目描述：

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

数据范围：

$3\le n \le 10^5, 1\le nums[i] \le 10^9$

题解：

关键在于如何求一个数组中大于 $val$ 的元素个数。如果是有序的，可以直接使用 upper_bound 解决。注意到 multiset是有序的，可以用这个，但是这个返回的是迭代器，求 $it, multiset.end$ 之间的元素个数，只能用 distance，这个复杂度是 $O(n)$ 的，太慢了。

考虑python中的 SortedList，需要导入包from sortedcontainers import SortedList，底层是对 list 的一些优化，返回的还是数组的下标。SortedList.bisect_right()

python中的二分查找，需要导入包 import bisect。bisect_right 查找 >，bisect_left查找 >=。

也有线段树的做法：

可以先对所有的数据排序，离散化，然后把离散化之后的数据作为数组下标，数组中记录出现的次数，这样的话，查询大于 val 的元素个数，只需要求 [val+1,n] 的后缀和。

考虑线段树单点修改，每次对 hash(nums[i]) 下标加 1。然后区间查询，查询区间的元素个数之和。

可以查询 [1, hash(nums[i])] 或者 [hash(nums[i]), n]。

也可以使用树状数组的做法：

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
import bisect
from sortedcontainers import SortedList


class Solution:
    def resultArray(self, nums):
        n = len(nums)
        arr1 = [nums[0]]
        mp1 = SortedList([nums[0]])
        arr2 = [nums[1]]
        mp2 = SortedList([nums[1]])
        for i in range(2, n):
            index1 = mp1.bisect_right(nums[i])
            index2 = mp2.bisect_right(nums[i])
            cnt1 = len(mp1) - index1
            cnt2 = len(mp2) - index2

            if cnt1 > cnt2:
                arr1.append(nums[i])
                mp1.add(nums[i])
            elif cnt1 < cnt2:
                arr2.append(nums[i])
                mp2.add(nums[i])
            else:
                if len(mp1) <= len(mp2):
                    arr1.append(nums[i])
                    mp1.add(nums[i])
                else:
                    arr2.append(nums[i])
                    mp2.add(nums[i])
        return arr1 + arr2
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
auto optimize_cpp_stdio = []()
{
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);
    return 0;
}();
struct SegTree
{
    struct TreeNode
    {
        int val;
        int sum;
        int l, r;
    };
    vector<TreeNode> tree;
    SegTree(int n)
    {
        tree.resize((n + 1) << 2);
    }
    void build(int cur, int l, int r)
    {
        tree[cur].l = l, tree[cur].r = r;
        tree[cur].sum = 0;
        tree[cur].val = 0;
        if (l == r)
            return;
        int mid = (l + r) >> 1;
        build(cur << 1, l, mid);
        build(cur << 1 | 1, mid + 1, r);
    }
    void update(int cur, int index, int k)
    {
        if (tree[cur].l == tree[cur].r)
        {
            tree[cur].val += k;
            tree[cur].sum += k;
            return;
        }
        int mid = tree[cur].l + tree[cur].r >> 1;
        if (index <= mid)
            update(cur << 1, index, k);
        if (index >= mid + 1)
            update(cur << 1 | 1, index, k);
        tree[cur].sum = tree[cur << 1].sum + tree[cur << 1 | 1].sum;
    }
    int query(int cur, int l, int r)
    {
        if (l > r)
            return 0;
        if (l <= tree[cur].l && tree[cur].r <= r)
            return tree[cur].sum;
        int mid = tree[cur].l + tree[cur].r >> 1;
        int sum = 0;
        if (l <= mid)
            sum += query(cur << 1, l, r);
        if (mid + 1 <= r)
            sum += query(cur << 1 | 1, l, r);
        return sum;
    }
};
class Solution
{
public:
    const static int maxn = 1e5 + 10;
    const static int maxm = 1e5 + 10;
    const static long long mod = 1e9 + 7;
    const long long INF_LL = 0x3f3f3f3f3f3f3f3f;
    const int INF = 0x3f3f3f3f;
    unordered_map<int, int> mp;
    int size = 0;
    int hash(int num)
    {
        if (!mp.count(num))
            mp[num] = ++size;
        return mp[num];
    }
    vector<int> resultArray(vector<int> &nums)
    {
        int n = nums.size();
        vector<int> arr1, arr2;
        vector<int> tmp = nums;
        sort(tmp.begin(), tmp.end());
        for (auto &x : tmp)
            hash(x);
        SegTree seg1(size), seg2(size);
        seg1.build(1, 1, size);
        seg2.build(1, 1, size);
        arr1.emplace_back(nums[0]);
        seg1.update(1, hash(nums[0]), 1);
        arr2.emplace_back(nums[1]);
        seg2.update(1, hash(nums[1]), 1);
        for (int i = 2; i < n; ++i)
        {
            int p = hash(nums[i]);
            int cnt1 = seg1.query(1, p + 1, size);
            int cnt2 = seg2.query(1, p + 1, size);
            if (cnt1 > cnt2)
            {
                arr1.emplace_back(nums[i]);
                seg1.update(1, p, 1);
            }
            else if (cnt1 < cnt2)
            {
                arr2.emplace_back(nums[i]);
                seg2.update(1, p, 1);
            }
            else
            {
                if (arr1.size() <= arr2.size())
                {
                    arr1.emplace_back(nums[i]);
                    seg1.update(1, p, 1);
                }
                else
                {
                    arr2.emplace_back(nums[i]);
                    seg2.update(1, p, 1);
                }
            }
        }
        arr1.insert(arr1.end(), arr2.begin(), arr2.end());
        return arr1;
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
auto optimize_cpp_stdio = []()
{
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);
    return 0;
}();
struct SegTree
{
    struct TreeNode
    {
        int val;
        int sum;
        int l, r;
    };
    vector<TreeNode> tree;
    SegTree(int n)
    {
        tree.resize((n + 1) << 2);
    }
    void build(int cur, int l, int r)
    {
        tree[cur].l = l, tree[cur].r = r;
        tree[cur].sum = 0;
        tree[cur].val = 0;
        if (l == r)
            return;
        int mid = (l + r) >> 1;
        build(cur << 1, l, mid);
        build(cur << 1 | 1, mid + 1, r);
    }
    void update(int cur, int index, int k)
    {
        if (tree[cur].l == tree[cur].r)
        {
            tree[cur].val += k;
            tree[cur].sum += k;
            return;
        }
        int mid = tree[cur].l + tree[cur].r >> 1;
        if (index <= mid)
            update(cur << 1, index, k);
        if (index >= mid + 1)
            update(cur << 1 | 1, index, k);
        tree[cur].sum = tree[cur << 1].sum + tree[cur << 1 | 1].sum;
    }
    int query(int cur, int l, int r)
    {
        if (l > r)
            return 0;
        if (l <= tree[cur].l && tree[cur].r <= r)
            return tree[cur].sum;
        int mid = tree[cur].l + tree[cur].r >> 1;
        int sum = 0;
        if (l <= mid)
            sum += query(cur << 1, l, r);
        if (mid + 1 <= r)
            sum += query(cur << 1 | 1, l, r);
        return sum;
    }
};
class Solution
{
public:
    const static int maxn = 1e5 + 10;
    const static int maxm = 1e5 + 10;
    const static long long mod = 1e9 + 7;
    const long long INF_LL = 0x3f3f3f3f3f3f3f3f;
    const int INF = 0x3f3f3f3f;
    unordered_map<int, int> mp;
    int size = 0;
    int hash(int num)
    {
        if (!mp.count(num))
            mp[num] = ++size;
        return mp[num];
    }
    vector<int> resultArray(vector<int> &nums)
    {
        int n = nums.size();
        vector<int> arr1, arr2;
        vector<int> tmp = nums;
        sort(tmp.begin(), tmp.end());
        for (auto &x : tmp)
            hash(x);
        SegTree seg1(size), seg2(size);
        seg1.build(1, 1, size);
        seg2.build(1, 1, size);
        arr1.emplace_back(nums[0]);
        seg1.update(1, hash(nums[0]), 1);
        arr2.emplace_back(nums[1]);
        seg2.update(1, hash(nums[1]), 1);
        for (int i = 2; i < n; ++i)
        {
            int p = hash(nums[i]);
            int cnt1 = seg1.query(1, p + 1, size);
            int cnt2 = seg2.query(1, p + 1, size);
            if (cnt1 > cnt2)
            {
                arr1.emplace_back(nums[i]);
                seg1.update(1, p, 1);
            }
            else if (cnt1 < cnt2)
            {
                arr2.emplace_back(nums[i]);
                seg2.update(1, p, 1);
            }
            else
            {
                if (arr1.size() <= arr2.size())
                {
                    arr1.emplace_back(nums[i]);
                    seg1.update(1, p, 1);
                }
                else
                {
                    arr2.emplace_back(nums[i]);
                    seg2.update(1, p, 1);
                }
            }
        }
        arr1.insert(arr1.end(), arr2.begin(), arr2.end());
        return arr1;
    }
};