在numpy数组中建立索引

Question

给定一个查询表：

colors = [  [0,0,0],\
            [0, 255, 0],\
            [0, 0, 255],\
            [255, 0, 0]]

索引numpy矩阵输入2x2：

a = np.array([[0,1],[1,1]])

如何将a映射到2x2x3矩阵b，其中b[i][j] = colors[a[i][j]] ？ 我想避免在这里使用for循环。

Answer 1

你有没有尝试过：

colors[a]

这是一个完整的示例：

import numpy as np

colors = np.array([[0,0,0],
                   [0, 255, 0],
                   [0, 0, 255],
                   [255, 0, 0]
                  ])
a = np.array([[0, 1], [1, 1]])

new = colors[a]
new.shape
# (2, 2, 3)
new
# array([[[  0,   0,   0],
#         [  0, 255,   0]],
#
#        [[  0, 255,   0],
#         [  0, 255,   0]]])