LeetCode 🔴 3072. Distribute Elements Into Two Arrays II

🔗 🔴 3072. Distribute Elements Into Two Arrays II 2053

tags: Weekly Contest 387 模擬(Simulation) 二分搜尋(Binary Search) 位運算(Bit Manipulation) 樹狀陣列(BIT) 線段樹(Segment Tree) 有序容器(Sorted Container)

題意

給定一個長度為 $n$ 的整數陣列 nums，下標從 $1$ 開始。

我們定義一個函數 greaterCount，使得 greaterCount(arr, val) 返回陣列 $arr$ 中 嚴格大於 $val$ 的元素數量。

你需要在 $n$ 次操作中將 $nums$ 的所有元素分配到兩個陣列 $arr1$ 和 $arr2$ 中。在第一次操作中，將 $nums[1]$ 附加到 $arr1$。在第二次操作中，將 $nums[2]$ 附加到 $arr2$。之後，在第 $i$ 次操作中：

• 如果 $\text{greaterCount}(arr1, nums[i]) > \text{greaterCount}(arr2, nums[i])$，則將 $nums[i]$ 添加到 $arr1$。
• 如果 $\text{greaterCount}(arr1, nums[i]) < \text{greaterCount}(arr2, nums[i])$，則將 $nums[i]$ 添加到 $arr2$。
• 如果 $\text{greaterCount}(arr1, nums[i]) == \text{greaterCount}(arr2, nums[i])$，則將 $nums[i]$ 添加到元素數量較少的陣列中。
• 如果仍然平手，則將 $nums[i]$ 添加到 $arr1$。

將 $arr1$ 和 $arr2$ 串聯形成陣列 $ans$ 後返回。

約束條件

• $3 \leq n \leq 10^5$
• $1 \leq nums[i] \leq 10^9$

思路

方法一：模擬(Simulation) + 二分搜尋(Binary Search)

若 $arr$ 為有序的，則可以利用 二分搜尋(Binary Search) 來計算 greaterCount。因此，我們可以將 $arr1$ 和 $arr2$ 作為有序陣列，並在每次操作中使用二分搜尋來計算 greaterCount。

• 令 $idx$ 為在 $arr$ 中第一個嚴格大於 $val$ 的元素的下標，則 $\text{greaterCount}(arr, val) = \text{len}(arr) - idx$。
• 此外，若需要將 $val$ 插入到 $arr$ 中，則 $idx$ 同時也為 $val$ 在 $arr$ 中應該插入的位置。
• 可以利用 bisect.bisect_right 或 upper_bound 來計算 $idx$。

具體步驟如下：

1. 初始化 $arr1$ 和 $arr2$，分別添加 $nums[0]$ 和 $nums[1]$。
2. 遍歷 $nums$ 中剩餘的元素，對於每一個元素，使用二分搜尋找到它在已排序的 $v1$ 和 $v2$ 中應該插入的位置 $idx1$ 和 $idx2$。
3. 根據插入位置計算出當前元素在 $v1$ 和 $v2$ 中分別有多少個元素大於它。
   • $gc1 = \text{len}(arr1) - idx1$
   • $gc2 = \text{len}(arr2) - idx2$
4. 比較這兩個數量，將元素添加到對應的陣列中。如果數量相等，則將元素添加到長度較短的陣列中。如果長度也相等，則將元素添加到 $arr1$。
   • 若$gc1 > gc2$ 或 $gc1 == gc2$ 且 $\text{len}(arr1) \leq \text{len}(arr2)$，則將元素添加到 $arr1$；
   • 否則，將元素添加到 $arr2$。
5. 最後將 $arr1$ 和 $arr2$ 串聯起來形成結果陣列 $ans$。

複雜度分析

• 時間複雜度：$\mathcal{O}(n^2)$
  • 每次二分搜尋的時間複雜度為 $\mathcal{O}(\log n)$，總共有 $n$ 次操作。
  • 插入到陣列中指定位置的時間複雜度為 $\mathcal{O}(n)$，總共有 $n$ 次操作。
• 空間複雜度：$\mathcal{O}(n)$

PythonC++

class Solution:
    def resultArray(self, nums: List[int]) -> List[int]:
        n = len(nums)
        arr1, arr2 = [nums[0]], [nums[1]] # ans
        v1, v2 = [nums[0]], [nums[1]] # sorted array

        for x in nums[2:]:
            idx1, idx2 = bisect_right(v1, x), bisect_right(v2, x)
            gc1, gc2 = len(arr1) - idx1, len(arr2) - idx2
            if gc1 > gc2 or gc1 == gc2 and len(arr1) <= len(arr2):
                arr1.append(x)
                v1.insert(idx1, x) # 用前面Binary Search的結果插入
            else:
                arr2.append(x)
                v2.insert(idx2, x)
        return arr1 + arr2

class Solution {
public:
    vector<int> resultArray(vector<int>& nums) {
        int n = nums.size();
        vector<int> arr1 = {nums[0]}, arr2 = {nums[1]}; // answer array
        vector<int> v1 = {nums[0]}, v2 = {nums[1]}; // sorted array
        for (int i = 2; i < n; i++) {
            int x = nums[i];
            int idx1 = upper_bound(v1.begin(), v1.end(), x) - v1.begin(), idx2 = upper_bound(v2.begin(), v2.end(), x) - v2.begin();
            int gc1 = arr1.size() - idx1, gc2 = arr2.size() - idx2;
            if (gc1 > gc2 || (gc1 == gc2 && arr1.size() <= arr2.size())) {
                arr1.push_back(x);
                v1.insert(v1.begin() + idx1, x);
            } else {
                arr2.push_back(x);
                v2.insert(v2.begin() + idx2, x);
            }
        }
        arr1.insert(arr1.end(), arr2.begin(), arr2.end());
        return arr1;
    }
};

方法一優化：有序容器(Sorted Container)

方法一雖然相對直覺，但在插入元素時需要對陣列進行移動，時間複雜度較高。若選擇一個合適的資料結構，可以在插入元素時降低時間複雜度。我們可以使用 Python 中的 sortedcontainers 模組中的 SortedList 來處理，其底層是一個類似於 B+ Tree 的資料結構，可以在 $O(\log N)$ 的時間內進行插入、刪除和查詢操作。

註： C++ 中的 pbds 也有類似的資料結構，應該也能用來解這道題，但我不會用，所以這裡只提供 Python 的解法。

複雜度分析

• 時間複雜度：$\mathcal{O}(n \log n)$
  • 每次二分搜尋的時間複雜度為 $\mathcal{O}(\log n)$，總共有 $n$ 次操作。
  • 插入到有序容器中指定位置的時間複雜度為 $\mathcal{O}(\log n)$，總共有 $n$ 次操作。
• 空間複雜度：$\mathcal{O}(n)$

from sortedcontainers import SortedList

class Solution:
    def resultArray(self, nums: List[int]) -> List[int]:
        sl1, sl2 = SortedList([nums[0]]), SortedList([nums[1]])
        res1, res2 = [nums[0]], [nums[1]]

        for x in nums[2:]:
            gc1 = len(sl1) - sl1.bisect_right(x)
            gc2 = len(sl2) - sl2.bisect_right(x)
            if gc1 > gc2 or gc1 == gc2 and len(sl1) <= len(sl2):
                sl1.add(x)
                res1.append(x)
            else:
                sl2.add(x)
                res2.append(x)
        return res1 + res2

方法二：離散化樹狀陣列(BIT)

可以將求 $\text{greaterCount}(arr, val)$ 視為求 $[0, val]$ 區間的元素個數，即前綴和。但我們每次都需要修改陣列，因此可以使用 樹狀陣列(BIT) 來維護這個區間的元素個數。

此外，由於 $nums[i]$ 的值域 $[1, 10^9]$ 過大，但注意到 $n < 10^5$，因此可以將其 離散化 為 $[0, m]$，其中 $m$ 為 $nums$ 中不重複元素的個數。離散化的過程如下：

• 將 $nums$ 中的元素去重後排序，得到 $sorted\_nums$。
• 對於每一個 $nums[i]$，可以用二分搜尋找到其在 $sorted\_nums$ 中的位置 $v$，則 $v$ 即為 $nums[i]$ 在離散化後的值。

之後便可以用兩個樹狀陣列來模擬 $arr1$ 和 $arr2$ ，分別為 $bit1$ 和 $bit2$，具體步驟如下：

1. 初始化 $arr1$ 和 $arr2$，分別添加 $nums[0]$ 和 $nums[1]$；同時初始化 $bit1$ 和 $bit2$ 為大小為 $m+1$ 的樹狀陣列，並分別在 $nums[0]$ 和 $nums[1]$ 離散化值的對應位置上加 $1$。
2. 遍歷 $nums$ 中剩餘的元素，對於每一個元素，使用二分搜尋找到它在 $sorted\_nums$ 中的位置 $v$，則 $v$ 即為 $nums[i]$ 的離散化值。
3. 根據 $bit1$ 和 $bit2$ 的前綴和來計算 $\text{greaterCount}(arr, val)$。
   • $\text{greaterCount}(arr, val) = \text{len}(arr) - \text{bit.preSum}(v)$
   • $\text{bit.preSum}(v)$ 表示 $\leq v$ 的元素個數，故 $\text{len}(arr) - \text{bit.preSum}(v)$ 表示 > v 的元素個數
4. 比較這兩個數量，將元素添加到對應的陣列中。如果數量相等，則將元素添加到長度較短的陣列中。如果長度也相等，則將元素添加到 $arr1$。
   • 若$gc1 > gc2$ 或 $gc1 == gc2$ 且 $\text{len}(arr1) \leq \text{len}(arr2)$，則將元素添加到 $arr1$，並在 $bit1$ 中對應位置上加 $1$；
   • 否則，將元素添加到 $arr2$，並在 $bit2$ 中對應位置上加 $1$。
5. 最後將 $arr1$ 和 $arr2$ 串聯起來形成結果陣列 $ans$。

複雜度分析

• 時間複雜度：$\mathcal{O}(n\log n)$
  • 離散化的時間複雜度為 $\mathcal{O}(n\log n)$；
  • 每次樹狀陣列的操作時間複雜度為 $\mathcal{O}(\log n)$，總共有 $n$ 次操作。
• 空間複雜度：$\mathcal{O}(n)$

PythonC++

class FenwickTree: # PURQ (Point Update Range Query), 0-based, initialization: O(nlogn)
    __slots__ = 'tree'

    def __init__(self, n: int): # 下標從 0 開始
        self.tree = [0] * n

    def add(self, k: int, x: int) -> None: # 令 nums[k] += x
        k += 1
        while k <= len(self.tree):
            self.tree[k - 1] += x
            k += (k & -k)

    def preSum(self, k: int) -> int: # 求前綴和: 求 nums[0] 到 nums[k] 的區間和
        res = 0
        k += 1
        while k > 0:
            res += self.tree[k - 1]
            k -= (k & -k)
        return res

class Solution:
    def resultArray(self, nums: List[int]) -> List[int]:
        n = len(nums)
        sorted_nums = sorted(set(nums)) # 離散化
        m = len(sorted_nums)
        arr1, arr2 = [nums[0]], [nums[1]]
        bit1, bit2 = FenwickTree(m + 1), FenwickTree(m + 1)
        bit1.add(bisect_left(sorted_nums, nums[0]), 1)
        bit2.add(bisect_left(sorted_nums, nums[1]), 1)
        for i in range(2, n):
            v = bisect_left(sorted_nums, nums[i])
            gc1 = len(arr1) - bit1.preSum(v) # bit1.preSum(v) 表示 <= v 的元素個數，故 len(arr1) - bit1.preSum(v) 表示 > v 的元素個數
            gc2 = len(arr2) - bit2.preSum(v)
            if gc1 > gc2 or (gc1 == gc2 and len(arr1) <= len(arr2)):
                arr1.append(nums[i])
                bit1.add(v, 1)
            else:
                arr2.append(nums[i])
                bit2.add(v, 1)
        return arr1 + arr2

class FenwickTree { // PURQ (Point Update Range Query), 0-based, initialization: O(nlogn)
private:
    int n = 0;
    vector<int> tree;
public:
    FenwickTree(int n) {
        this->n = n;
        tree = vector<int>(n, 0);
    }
    void add(int k, int x) { // 令 nums[k] += x
        for (int i = k+1; i <= n; i += i & -i) {
            tree[i-1] += x;
        }
    }
    int preSum(int k) { // 求前綴和: 求 nums[0] 到 nums[k] 的區間和
        int res = 0;
        for (int i = k+1; i > 0; i -= i & -i) {
            res += tree[i-1];
        }
        return res;
    }
};

class Solution {
public:
    vector<int> resultArray(vector<int>& nums) {
        int n = nums.size();
        vector<int> sorted_nums = nums; // 離散化
        sort(sorted_nums.begin(), sorted_nums.end());
        sorted_nums.erase(unique(sorted_nums.begin(), sorted_nums.end()), sorted_nums.end()); // 去重
        int m = sorted_nums.size();
        vector<int> arr1 = {nums[0]}, arr2 = {nums[1]};
        FenwickTree bit1(m+1), bit2(m+1);
        int idx1 = lower_bound(sorted_nums.begin(), sorted_nums.end(), nums[0]) - sorted_nums.begin();
        int idx2 = lower_bound(sorted_nums.begin(), sorted_nums.end(), nums[1]) - sorted_nums.begin();
        bit1.add(idx1, 1);
        bit2.add(idx2, 1);
        for (int i = 2; i < n; i++) {
            int idx = lower_bound(sorted_nums.begin(), sorted_nums.end(), nums[i]) - sorted_nums.begin();
            int gc1 = arr1.size() - bit1.preSum(idx); // bit1.preSum(v) 表示 <= v 的元素個數，故 len(arr1) - bit1.preSum(v) 表示 > v 的元素個數
            int gc2 = arr2.size() - bit2.preSum(idx);
            if (gc1 > gc2 || (gc1 == gc2 && arr1.size() <= arr2.size())) {
                arr1.push_back(nums[i]);
                bit1.add(idx, 1);
            } else {
                arr2.push_back(nums[i]);
                bit2.add(idx, 1);
            }
        }
        arr1.insert(arr1.end(), arr2.begin(), arr2.end());
        return arr1;
    }
};

---

寫在最後

Cover photo is generated by @たろたろ, thanks for their work!