优化生成和拆分列表的内存使用

Question

因此，任务是生成从 0 到 N 的所有唯一数字对，这使得列表的长度为 N * (N - 1)。

为了生成所有唯一的对，我写了这个函数：

def generate(N):
    for i in range(N):
        for j in range(i + 1, N):
            yield(i, j)

接下来我应该做的是将列表 M 分成相等的大小，我已经写了这个：

def divide(full_list, N, M):
    l, r = divmod(N * (N - 1) // 2, M)
    return [full_list[i * l + min(i, r): (i + 1) * l + min(i + 1, r)] for i in range(M)]

拆分列表的标准正是它应该的样子。这一切都很好，但几乎没有限制：

N - 可以高达 ~100000

内存使用限制 - 256mb

时间限制 - 2s

当 N 很大时，上面的解决方案使用超过 256MB（不幸的是，我不知道它在哪个数量上失败）。

我该如何改进和优化它？

Answer 1

您可以将combinations 与经典的grouper 迭代器一起使用。这样你就有了一个纯流式迭代器（不需要存储长列表）：

# example
n = 4
m = 2

batch_size = n * (n - 1) // (2 * m)
[list(pairs) for pairs in grouper(combinations(range(n), 2), batch_size)]

# out:
[[(0, 1), (0, 2), (0, 3)], [(1, 2), (1, 3), (2, 3)]]

其他例子：

batch_size = 8
n = 256  # full list of pairs would be 32_640 long
first_k = 10

for batch in islice(grouper(combinations(range(n), 2), batch_size), first_k):
    print(list(batch))

输出：

[(0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), (0, 8)]
[(0, 9), (0, 10), (0, 11), (0, 12), (0, 13), (0, 14), (0, 15), (0, 16)]
[(0, 17), (0, 18), (0, 19), (0, 20), (0, 21), (0, 22), (0, 23), (0, 24)]
[(0, 25), (0, 26), (0, 27), (0, 28), (0, 29), (0, 30), (0, 31), (0, 32)]
[(0, 33), (0, 34), (0, 35), (0, 36), (0, 37), (0, 38), (0, 39), (0, 40)]
[(0, 41), (0, 42), (0, 43), (0, 44), (0, 45), (0, 46), (0, 47), (0, 48)]
[(0, 49), (0, 50), (0, 51), (0, 52), (0, 53), (0, 54), (0, 55), (0, 56)]
[(0, 57), (0, 58), (0, 59), (0, 60), (0, 61), (0, 62), (0, 63), (0, 64)]
[(0, 65), (0, 66), (0, 67), (0, 68), (0, 69), (0, 70), (0, 71), (0, 72)]
[(0, 73), (0, 74), (0, 75), (0, 76), (0, 77), (0, 78), (0, 79), (0, 80)]

更新：有余数的示例（n*(n-1)/2 不是m 的倍数）：

n = 5
m = 3

batch_size = n * (n - 1) // (2 * m)
[list(pairs) for pairs in grouper(combinations(range(n), 2), batch_size)]
# out:
[[(0, 1), (0, 2), (0, 3)],
 [(0, 4), (1, 2), (1, 3)],
 [(1, 4), (2, 3), (2, 4)],
 [(3, 4), None, None]]

如果有人喜欢删除这些 None 元素（grouper() 的默认 fillvalue），从而接受最后一批更小的：

[[p for p in pairs if p] for pairs in grouper(combinations(range(n), 2), batch_size)]
# out:
[[(0, 1), (0, 2), (0, 3)],
 [(0, 4), (1, 2), (1, 3)],
 [(1, 4), (2, 3), (2, 4)],
 [(3, 4)]]

更新 2：需要更长的初始批次。

似乎（没有太多明确的动机，但基于 OP 的 divide() 所做的）在有剩余的情况下（而不是 @987654336 @ 填充物，或最后一批更短的批次）。

不是问题，但与其更改漂亮、规范的 grouper，我们可以自己计算并拆分初始批次。请注意，这不会影响内存：整个设置仍然是一个纯流式迭代器：

def gen_pair_batch(n, m):
    n_pairs = n * (n - 1) // 2
    batch_size = n_pairs // m
    initial_batch = batch_size + n_pairs % batch_size
    
    gen_pairs = combinations(range(n), 2)
    yield tuple(islice(gen_pairs, initial_batch))
    yield from grouper(gen_pairs, batch_size)

例子：

>>> list(gen_pair_batch(5, 3))
[((0, 1), (0, 2), (0, 3), (0, 4)),
 ((1, 2), (1, 3), (1, 4)),
 ((2, 3), (2, 4), (3, 4))]

性能

上述任何一个版本都使用非常有限的内存量（这就是流迭代器的美妙之处）。但是，我怀疑您能否实现n=100_000 的最大时间 2s 的目标：

def go(n):
    n_pairs = n * (n - 1) // 2
    m = max(10, n_pairs // 1000)
    process = psutil.Process()
    mmin, mmax = sys.maxsize, 0
    count = 0
    for batch in gen_pair_batch(n, m):
        # ... do something with the batch of pairs
        mem = process.memory_info().rss
        if mem < mmin: mmin = mem
        if mem > mmax: mmax = mem
        count += 1

    mem1 = process.memory_info().rss
    return count, mmin, mmax

测试：

%%time
count, mmin, mmax = go(n)
print(f'finished go({n}) count = {count}; total resident mem was between: {mmin}, {mmax}')

对于n = 1000：

finished go(1000) count = 499; total resident mem was between: 124968960, 124968960
CPU times: user 20.9 ms, sys: 94 µs, total: 21 ms
Wall time: 20.1 ms

对于n = 10_000：

finished go(10000) count = 49995; total resident mem was between: 124968960, 124968960
CPU times: user 1.77 s, sys: 164 ms, total: 1.94 s
Wall time: 1.94 s

对于n = 100_000：

finished go(100000) count = 4999950; total resident mem was between: 124968960, 124968960
CPU times: user 2min 57s, sys: 17 s, total: 3min 14s
Wall time: 3min 14s

所以：没有时间限制，但好消息是内存使用情况良好。顺便说一句，我想知道您为什么要生成这些对（最多 n = 100K，以及为什么要批量生成等）。如果对实际/更大的目标有更深入的了解，我们可能会想出更有效的方法.