Numpy 的 ndindex() 适用于您给出的示例,但它并不适用于所有用例。与 Python 的内置 range() 不同,后者允许任意 start、stop 和 step,numpy 的 np.ndindex() 只接受 stop。 (start 假定为(0,0,...),step 为(1,1,...)。)
这是一个更像内置 range() 函数的实现。也就是说,它允许任意 start/stop/step 参数,但它适用于 元组 而不仅仅是整数。
import sys
from itertools import product, starmap
# Python 2/3 compatibility
if sys.version_info.major < 3:
from itertools import izip
else:
izip = zip
xrange = range
def ndrange(start, stop=None, step=None):
if stop is None:
stop = start
start = (0,)*len(stop)
if step is None:
step = (1,)*len(stop)
assert len(start) == len(stop) == len(step)
for index in product(*starmap(xrange, izip(start, stop, step))):
yield index
例子:
In [7]: for index in ndrange((1,2,3), (10,20,30), step=(5,10,15)):
...: print(index)
...:
(1, 2, 3)
(1, 2, 18)
(1, 12, 3)
(1, 12, 18)
(6, 2, 3)
(6, 2, 18)
(6, 12, 3)
(6, 12, 18)