我在哪里可以改进我的代码以缩短其执行时间？

Question

请求from HackerRank：

如果客户在某一天的消费金额大于或等于客户在过去几天的平均支出的 2 倍，他们会向客户发送有关潜在欺诈的通知。银行不会向客户发送任何通知，直到他们至少拥有前几天的交易数据。

给定过去的天数 d 和客户在 n 天期间的每日总支出，确定客户将收到通知的次数n 天。

我的代码可以解决问题，但是对于大型测试用例有时间限制。我的代码无法通过时间限制要求。我的代码实际上很短：

from statistics import median

first_multiple_input = input().rstrip().split()
n = int(first_multiple_input[0])
d = int(first_multiple_input[1])
expenditure = list(map(int, input().rstrip().split()))
count=0
for i in range(len(expenditure)-d):
    if expenditure[d+i] >= 2*median(expenditure[i:d+i]) :
        count+=1
print( count)

请指出造成延迟的原因以及如何改进。

帮助理解代码的小测试用例：

9 5                 expenditure[] size n =9, d = 5
2 3 4 2 3 6 8 4 5   expenditure = [2, 3, 4, 2, 3, 6, 8, 4, 5]

Answer 1

分析/想法

你的median(expenditure[i:d+i]) 是罪魁祸首，因为每次sorting 整个未排序的大小为 d 的切片都需要 O(d log d) 时间。您可以通过将尾随元素的当前窗口保留在 SortedList 中来将其减少到 O(log d)。您可以从中间元素或两个元素中获取中位数，然后更新，只需添加一个新元素并删除最旧的元素。

实施

from sortedcontainers import SortedList

n = 9
d = 5
expenditure = [2, 3, 4, 2, 3, 6, 8, 4, 5]

count = 0
trailing = SortedList(expenditure[:d])
half = d // 2
for i in range(d, n):
    median = (trailing[half] + trailing[~half]) / 2
    if expenditure[i] >= 2 * median:
        count += 1
    trailing.add(expenditure[i])
    trailing.remove(expenditure[i - d])
print(count)

我们可以省略/ 2 和2 *，但是“中位数”将是错误的名称，而naming things is hard。我们可以if expenditure[i] >= trailing[half] + trailing[~half]，但我觉得不太清楚。

输出

如果你添加

    print(f'{trailing=} {median=} {expenditure[i]=}')

在median = ... 行之后，您可以看到发生了什么：

trailing=SortedList([2, 2, 3, 3, 4]) median=3.0 expenditure[i]=6
trailing=SortedList([2, 3, 3, 4, 6]) median=3.0 expenditure[i]=8
trailing=SortedList([2, 3, 4, 6, 8]) median=4.0 expenditure[i]=4
trailing=SortedList([2, 3, 4, 6, 8]) median=4.0 expenditure[i]=5
2

替代实现

使用zip 代替索引：

count = 0
trailing = SortedList(expenditure[:d])
half = d // 2
for today, oldest in zip(expenditure[d:], expenditure):
    median = (trailing[half] + trailing[~half]) / 2
    if today >= 2 * median:
        count += 1
    trailing.add(today)
    trailing.remove(oldest)
print(count)

另一种数据结构：有序的正则列表

我发现了问题at HackerRank，它没有sortedcontainers。但以下内容在那里被接受。

我们可以使用常规的 Python list，但在 Python 标准库中包含的 sorted 和 bisect 的帮助下自行排序：

from bisect import bisect_left, insort

count = 0
trailing = sorted(expenditure[:d])
half = d // 2
for today, oldest in zip(expenditure[d:], expenditure):
    median = (trailing[half] + trailing[~half]) / 2
    if today >= 2 * median:
        count += 1
    insort(trailing, today)
    del trailing[bisect_left(trailing, oldest)]
print(count)

访问中间元素需要 O(1) 时间，查找插入/删除索引需要 O(log d) 时间，实际插入/删除需要 O(d) 时间（因为它需要移位索引右侧的所有元素）。但是这种 O(d) 转换非常快非常低。

二更：有序字节数组和计数排序

问题最初不包括对 HackerRank 的引用。现在我看到这些值仅限于 0 到 200 之间的整数，我们也可以使用 bytearray：

trailing = bytearray(sorted(expenditure[:d]))

正如我刚才在讨论中所指出的，对于这个允许值范围，我们还可以使用计数排序的形式。我认为Fenwick tree 会使这个速度特别快，我以后可能会尝试。

基准测试

在您提到的 cmets 中，n=200000 和 d=10122 是一个大案例。所以我用这个数据进行了测试：

n = 200000
d = 10122
expenditure = random.choices(range(201), k=n)

我的解决方案的基准：

                       at replit.com   on my weak laptop
SortedList + indexing   ~1.8 seconds    ~6.4 seconds
SortedList + zipping    ~1.8 seconds    ~6.4 seconds
sorted regular list     ~0.6 seconds    ~8.8 seconds
sorted bytearray        ~0.3 seconds    ~1.7 seconds

不知道为什么常规列表解决方案在我的笔记本电脑上相对较慢。我怀疑它超出了我 CPU 的 1 级缓存。