使用索引位置填充2维numpy数组

Question

我一直试图找到一种干净的、pythonic的方法来填充一个空的numpy数组中的每个元素，使该元素的索引值成为其值，而不使用for循环。对于1-D数组，可以很容易地使用类似np.arange或基本的range函数。但在2-D及更高维度的情况下，我不知道如何轻松地做到这一点。

（编辑：或者就像这样构建一个普通的列表，然后使用np.array(lst)将其变为数组。我想我刚刚回答了自己的问题——使用列表推导？）

例子：

rows = 4
cols = 4
arr = np.empty((rows, cols, 2))  # 4x4 matrix with [x,y] location

for y in range(rows):
    for x in range(cols):
        arr[y, x] = [y, x]

'''
Expected output:
[[[0,0], [0,1], [0,2], [0,3]],
 [[1,0], [1,1], [1,2], [1,3]],
 [[2,0], [2,1], [2,2], [2,3]],
 [[3,0], [3,1], [3,2], [3,3]]]
'''

- stormont

Answer 1

您展示的是一个4X4矩阵的网格，您可以使用np.mgrid，然后转置结果：

np.moveaxis(np.mgrid[:rows,:cols], 0, -1)
#array([[[0, 0],
#        [0, 1],
#        [0, 2],
#        [0, 3]],

#       [[1, 0],
#        [1, 1],
#        [1, 2],
#        [1, 3]],

#       [[2, 0],
#        [2, 1],
#        [2, 2],
#        [2, 3]],

#       [[3, 0],
#        [3, 1],
#        [3, 2],
#        [3, 3]]])

或者使用带有矩阵索引ij的np.meshgrid：

np.dstack(np.meshgrid(np.arange(rows), np.arange(cols), indexing='ij'))
#array([[[0, 0],
#        [0, 1],
#        [0, 2],
#        [0, 3]],

#       [[1, 0],
#        [1, 1],
#        [1, 2],
#        [1, 3]],

#       [[2, 0],
#        [2, 1],
#        [2, 2],
#        [2, 3]],

#       [[3, 0],
#        [3, 1],
#        [3, 2],
#        [3, 3]]])

- Psidom

---

这是我目前见过的最简洁的np.mgrid示例。文档对我来说一团糟，我也无法从其他SO帖子中理解。 - stormont

这似乎很容易扩展到更高的维度，这很好。我还没有弄清楚如何正确地转置3D所需的参数，但至少我可以告诉它是受支持的。 - stormont

总的来说，您可以使用 np.moveaxis（它是一个特殊的 transpose）而不是在此处使用 transpose。它应该适用于 nd，而无需修改 args。 - Psidom

如果需要将值作为字符串，怎么办？例如 arr[y, x] = "y_x" - lesolorzanov

Answer 2

另一种方法是使用np.indices和concatenate函数

 np.concatenate([x.reshape(4,4,1) for x in np.indices((4,4))],2)

或者使用np.dstack。

np.dstack(np.indices((4,4)))

---

由于存在很多可能性，因此需要进行一些基准测试。

def Psidom_mrgid(rows,cols):
    np.mgrid[:rows, :cols].transpose((1, 2, 0))

def Psidom_mesh(rows,cols):
    np.dstack(np.meshgrid(np.arange(rows), np.arange(cols), indexing='ij'))

def Mad_tile(rows,cols):
    r = np.tile(np.arange(rows).reshape(rows, 1), (1, cols))
    c = np.tile(np.arange(cols), (rows, 1))
    result = np.stack((r, c), axis=-1)

def bora_comp(rows,cols):
    x = [[[i, j] for j in range(rows)] for i in range(cols)]

def djk_ind(rows,cols):
    np.concatenate([x.reshape(rows, cols, 1) for x in np.indices((rows, cols))], 2)

def devdev_mgrid(rows,cols):
    index_tuple = np.mgrid[0:rows, 0:cols]
    np.dstack(index_tuple).reshape((rows, cols, 2)


In[8]: %timeit Psidom_mrgid(1000,1000)
100 loops, best of 3: 15 ms per loop

In[9]: %timeit Psidom_mesh(1000,1000)
100 loops, best of 3: 9.98 ms per loop

In[10]: %timeit Mad_tile(1000,1000)
100 loops, best of 3: 15.3 ms per loop

In[11]: %timeit bora_comp(1000,1000)
1 loop, best of 3: 221 ms per loop

In[12]: %timeit djk_ind(1000,1000)
100 loops, best of 3: 9.72 ms per loop

In[13]: %timeit devdev_mgrid(1000,1000)
10 loops, best of 3: 20.6 ms per loop

- DJK

Answer 3

我猜这很Pythonic：

[[[i,j] for j in range(5)] for i in range(5)]

输出:

[[[0, 0], [0, 1], [0, 2], [0, 3], [0, 4]],

[[1, 0], [1, 1], [1, 2], [1, 3], [1, 4]],

[[2, 0], [2, 1], [2, 2], [2, 3], [2, 4]],

[[3, 0], [3, 1], [3, 2], [3, 3], [3, 4]],

[[4, 0], [4, 1], [4, 2], [4, 3], [4, 4]]]

- Bora

Answer 4

查看 numpy.mgrid，它将返回带有 i 和 j 索引的两个数组。 要将它们组合起来，您可以堆叠这些数组并调整它们的形状。类似于这样：

import numpy as np

def index_pair_array(rows, cols):
    index_tuple = np.mgrid[0:rows, 0:cols]
    return np.dstack(index_tuple).reshape((rows, cols, 2))

- Devin Cairns

---

@Psidom的回答更好，您不需要堆叠np.mgrid输出，因为它已经是一个ndarray对象。 - Devin Cairns

Answer 5

有几种numpy式的方法可以做到这一点。

其中一种方法是使用np.tile和np.stack：

r = np.tile(np.arange(rows).reshape(rows, 1), (1, cols))
c = np.tile(np.arange(cols), (rows, 1))
result = np.stack((r, c), axis=-1)

更好的获取坐标的方法可能是使用np.meshgrid：

rc = np.meshgrid(np.arange(rows), np.arange(cols), indexing='ij')
result = np.stack(rc, axis=-1)

- Mad Physicist