iT邦幫忙

1

發現一個python好用函式庫,能夠維護排序的陣列- Leetcode 1649. Create Sorted Array through Instructions

今天在刷leetcode時看到一道有意思的題,
話不多說先上題目
1649. Create Sorted Array through Instructions

題意: 給定陣列instructions,按順序將元素插入空陣列nums中,
使nums成為排序陣列
(其實就是insertion sort)

假設插入元素n時,
陣列nums比n小的元素有a個,
陣列nums比n大的元素有b個,
每次插入的cost定義為min(a,b)
計算插入所有元素所需的cost
(由於答案可能很大,答案請回傳10^9+7)

範例:
Input: instructions = [1,2,3,6,5,4]
Output: 3
Explanation: Begin with nums = [].
Insert 1 with cost min(0, 0) = 0, now nums = [1].
Insert 2 with cost min(1, 0) = 0, now nums = [1,2].
Insert 3 with cost min(2, 0) = 0, now nums = [1,2,3].
Insert 6 with cost min(3, 0) = 0, now nums = [1,2,3,6].
Insert 5 with cost min(3, 1) = 1, now nums = [1,2,3,5,6].
Insert 4 with cost min(3, 2) = 2, now nums = [1,2,3,4,5,6].
The total cost is 0 + 0 + 0 + 0 + 1 + 2 = 3.

測資限制:

  • 1 <= instructions.length <= 10^5

這一要我們計算類似insertion sort的方法的代價,
由於題目測資還蠻大的,陣列最大可以到十萬,
insertion sort的時間複雜度又是O(n^2),
如果直接用insertion sort來模擬題目操作的話肯定超時

比較容易想到的維護一個排序陣列,用「二分搜索」來找每一次插入元素時,
比該元素小、大的元素各有幾個,
在排序陣列中搜索比自己小的元素個數只要O(log n)的時間,
然後再把該元素放入排序陣列中

思路如底下程式碼所示:

each search O(log n), insert O(n)- 超時

from bisect import bisect_left, bisect_right
class Solution:
    def createSortedArray(self, instructions: List[int]) -> int:
        arr = []
        cost = 0
        for e in instructions:
            left = bisect_left(arr, e)
            right = len(arr)-bisect_right(arr, e)
            cost += min(left, right)
            arr.insert(left,e)
        return cost%(10**9+7)

https://ithelp.ithome.com.tw/upload/images/20210411/201359198mZzlmk6Qg.png

程式碼內的left,right是用二分搜索計算插入某個元素比它小、大的元素各有幾個。
很可惜程式碼還是超時了,
問題在於雖然二分搜索可以做到O(log n)查詢,
但是把一個新的元素放進排序陣列卻需要O(n),
放了n個元素整體時間還是O(n^2)

那麼有沒有能夠讓插入做的更快的結構呢?
有種資料結構叫作binary search tree,
用一些技巧可以做到插入,查找只要O(log n)的時間,
但我相信用binary search tree來解的話學習門檻是比較高的

於是本文的逆天好用工具SortedList登場

SortedList簡介

SortedList來自函式庫sortedcontainers,
雖然不是python的標準內建函式庫,
但於leetcode實測是可以使用的,
它的效果是可以維護一個排序陣列,
做到插入,查找只要O(log n)的時間,
而且SortedList的API與內建模組bisect是相近的,

有了SortedList這項可以維護排序陣列的工具,
這樣便無需手刻一個binary search tree出來,
功能亦比bisect模組更加齊全

感謝python,總是有很實用的函式庫,
讓學習python類似拼樂高積木般,
能夠用簡單的程式碼實作實用的功能
上述程式碼只要修改一點即可把這道難題給解了

函式庫SortedList,each search&insert O(log n)

from sortedcontainers import SortedList
class Solution:
    def createSortedArray(self, instructions: List[int]) -> int:
        arr = SortedList()
        cost = 0
        for e in instructions:
            left = arr.bisect_left(e)
            right = len(arr)-arr.bisect_right(e)
            cost += min(left, right)
            arr.add(e) #在排序陣列中添加新元素,複雜度O(log n)
        return cost%(10**9+7)

https://ithelp.ithome.com.tw/upload/images/20210411/20135919wj5m43FWpM.png

函式庫文檔

對於上述提到的bisectSortedList函式庫,
有興趣的同學可以google一下用法


尚未有邦友留言

立即登入留言