python 一个numpy函数,列出了nd-array中所有可能的“坐标”

6l7fqoea  于 2023-04-28  发布在  Python
关注(0)|答案(1)|浏览(75)

是否有一个numpy函数可以列出numpy.ndarray中所有可能的“坐标”?

它会做我的自定义函数all_possible_coordinates所做的事情:

import numpy as np

def all_possible_coordinates(shape): 
    return np.array([item.ravel() for item in np.meshgrid(*(np.arange(n) for n in shape))]).T

def test_all_possible_coordinates():
    shape = (1, 2, 2)
    assert np.all(all_possible_coordinates(shape)
                  == np.array([[0, 0, 0],
                               [0, 0, 1],
                               [0, 1, 0],
                               [0, 1, 1]]))

test_all_possible_coordinates()

编辑:更优雅,由M Z

def all_possible_coordinates(shape): 
    return np.argwhere(np.ones(shape))
os8fio9y

os8fio9y1#

可以使用numpy.ndindex

In [1]: import numpy as np

In [2]: list(np.ndindex(1, 2, 2))
Out[2]: [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1)]

In [3]: for idx in np.ndindex(2, 3, 2):
   ...:     print(idx)
   ...: 
(0, 0, 0)
(0, 0, 1)
(0, 1, 0)
(0, 1, 1)
(0, 2, 0)
(0, 2, 1)
(1, 0, 0)
(1, 0, 1)
(1, 1, 0)
(1, 1, 1)
(1, 2, 0)
(1, 2, 1)

相关问题