python - 如何使用numpy生成楼梯编号序列

Question

我正在尝试使用 python 生成数字序列列表，如下所示。

[0,0,0,0] [0,0,0,1] [0,0,0,2] [0,0,1,1] [0,0,1,2] [0,0,2,2] [0,1,1,1]
[0,1,1,2] [0,1,2,2] [0,2,2,2] [1,1,1,1] [1,1,1,2] ... [2,2,2,2]

现在，我可以使用带有递归调用的纯 Python 来完成此操作，但单次运行需要花费大量时间(几个小时)。我想知道是否可以使用 numpy 来做到这一点并节省大量时间，如果可以，怎么做？

Answer 1

您要找的是itertools.combinations_with_replacement .来自文档:

combinations_with_replacement('ABCD', 2)
AA AB AC AD BB BC BD CC CD DD

因此:

>>> import itertools as it
>>> list(it.combinations_with_replacement((0, 1, 2), 4))
[(0, 0, 0, 0), (0, 0, 0, 1), (0, 0, 0, 2),
 (0, 0, 1, 1), (0, 0, 1, 2), (0, 0, 2, 2),
 (0, 1, 1, 1), (0, 1, 1, 2), (0, 1, 2, 2),
 (0, 2, 2, 2), (1, 1, 1, 1), (1, 1, 1, 2),
 (1, 1, 2, 2), (1, 2, 2, 2), (2, 2, 2, 2)]

这个方法最好的部分是因为它返回一个生成器，你可以在不存储它的情况下迭代它。这是一个巨大的优势，因为它会为您节省大量内存。

其他实现和时序

这里有一些更多的实现，包括一个 numpy 实现。 combinations_with_replacement(try2 函数)似乎是最快的:

import itertools as it
import timeit

import numpy as np

def try1(n, m):
    return [t for t in it.product(range(n), repeat=m) if all(a <= b for a, b in zip(t[:-1], t[1:]))]

def try2(n, m):
    return list(it.combinations_with_replacement(range(n), m))

def try3(n, m):
    a = np.mgrid[(slice(0, n),) * m] # All points in a 3D grid within the given ranges
    a = np.rollaxis(a, 0, m + 1)     # Make the 0th axis into the last axis
    a = a.reshape((-1, m))           # Now you can safely reshape while preserving order
    return a[np.all(a[:, :-1] <= a[:, 1:], axis=1)]

>>> %timeit b = try1(3, 4)
10000 loops, best of 3: 78.1 µs per loop
>>> %timeit b = try2(3, 4)
The slowest run took 8.04 times longer than the fastest. This could mean that an intermediate result is being cached.
1000000 loops, best of 3: 1.66 µs per loop
>>> %timeit b = try3(3, 4)
10000 loops, best of 3: 97.8 µs per loop

即使对于更大的数字也是如此:

>>> %timeit b = try1(3, 6)
1000 loops, best of 3: 654 µs per loop
>>> %timeit b = try2(3, 6)
The slowest run took 7.20 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 3: 2.33 µs per loop
>>> %timeit b = try3(3, 6)
10000 loops, best of 3: 166 µs per loop

注意事项:

• 我使用的是 python3
• 我用了this answer用于try1的实现。
• 我用了this answer用于try3的实现。