高效的 Python 切片赋值

 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

from bisect import bisect_left
from copy import deepcopy
from random import randint
from timeit import timeit
from typing import Callable, List

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

sns.set_style("whitegrid")
plt.rcParams["figure.figsize"] = (12, 8)


def generate_samples(row: int, col: int) -> List[List[int]]:
    return [[randint(1, 100) for _ in range(col)] for _ in range(row)]


def test1(arr: List[int]) -> None:
    sorted_arr = []
    for x in arr:
        k = bisect_left(sorted_arr, x)
        sorted_arr[k:k] = [x]
    # assert all(b >= a for a, b in zip(sorted_arr[:-1], sorted_arr[1:]))


def test2(arr: List[int]) -> None:
    sorted_arr = []
    for x in arr:
        k = bisect_left(sorted_arr, x)
        sorted_arr.insert(k, x)
    # assert all(b >= a for a, b in zip(sorted_arr[:-1], sorted_arr[1:]))


def run(test: Callable[[List[int]], None], samples) -> List[float]:
    res = []
    for sample in samples:
        res.append(timeit(lambda: test(sample), number=10))
    return res

ndim = 100000
num_samples = 10
data = generate_samples(num_samples, ndim)
data_copy1 = deepcopy(data)
data_copy2 = deepcopy(data)

res1 = run(test1, data_copy1)
res2 = run(test2, data_copy2)

res = []
labels = []

for a, b in zip(res1, res2):
    res.append(a)
    res.append(b)
    labels.append("test1")
    labels.append("test2")

data = pd.DataFrame(data={"Data": list(range(len(res))), "Time": res, "Label": labels})
sns.barplot(x="Data", y="Time", hue="Label", data=data, dodge=False)
plt.savefig("result.png", bbox_inches="tight")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

class MedianFinder:
    def __init__(self):
        self.S = []

    def addNum(self, num: int) -> None:
        k = bisect_left(self.S, num)
        self.S[k:k] = [num]

    def findMedian(self) -> float:
        n = len(self.S)
        return (self.S[n // 2] + self.S[(n - 1) // 2]) / 2

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        n = len(nums)
        ans = 0
        S = [0]
        cumsum = 0
        for x in nums:
            cumsum += x
            ans += bisect_right(S, cumsum - lower) - bisect_left(S, cumsum - upper)
            k = bisect_left(S, cumsum)
            S[k:k] = [cumsum]
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        ans, S, n = [], [], len(nums)
        for i in range(n):
            j = bisect_left(S, nums[i])
            S[j:j] = [nums[i]]
            if len(S) == k:
                ans.append((S[k // 2] + S[(k - 1) // 2]) / 2)
                j = bisect_left(S, nums[i - k + 1])
                S[j : j + 1] = []
        return ans

1
2
3
4
5
6
7
8
9

class Solution:
    def reversePairs(self, nums: List[int]) -> int:
        ans, S = 0, []
        for x in nums:
            k = bisect_right(S, 2 * x)
            ans += len(S) - k
            k = bisect_left(S, x)
            S[k:k] = [x]
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

class Solution:
    def closestRoom(self, rooms: List[List[int]], queries: List[List[int]]) -> List[int]:
        n = len(queries)
        ans = [-1] * n
        index = list(range(n))
        S = []

        index.sort(key=lambda i: queries[i][1], reverse=True)
        rooms.sort(key=lambda x: x[1])

        for i in index:
            preferred, min_size = queries[i]
            while rooms and rooms[-1][1] >= min_size:
                k = bisect_left(S, rooms[-1][0])
                S[k:k] = [rooms[-1][0]]
                rooms.pop()
            if S:
                k = bisect_left(S, preferred)
                if k != len(S):
                    ans[i] = S[k]
                if k != 0:
                    if ans[i] == -1 or abs(S[k - 1] - preferred) <= abs(ans[i] - preferred):
                        ans[i] = S[k - 1]
        return ans