根据数组给出的索引选择矩阵上的项目

Selecting items on a matrix based on indexes given by an array

考虑这个矩阵:

[0.9, 0.45, 0.4, 0.35],
[0.4, 0.8, 0.3, 0.25],
[0.5, 0.45, 0.9, 0.35],
[0.2, 0.18, 0.8, 0.1],
[0.6, 0.45, 0.4, 0.9]

和这个列表:

[0,1,2,3,3]

我想创建一个如下所示的列表:

[0.9, 0.8, 0.9, 0.1, 0.9]

澄清一下,对于每一行,我想要其列索引包含在第一个数组中的矩阵元素。我怎样才能做到这一点?

将两个列表压缩在一起,如下所示

a=[[0.9, 0.45, 0.4, 0.35],[0.4, 0.8, 0.3, 0.25],[0.5, 0.45, 0.9, 0.35],[0.2, 0.18, 0.8, 0.1],[0.6, 0.45, 0.4, 0.9]]
b=[0,1,2,3,3]
[i[j] for i,j in zip(a,b)]

结果

[0.9, 0.8, 0.9, 0.1, 0.9]

这基本上将矩阵中的每个子列表与第二个列表的元素按 zip(a,b)

的顺序配对

然后为每一对选择 a

的第 b 个元素

如果这是一个 numpy 数组,您可以传入两个 numpy 数组来访问所需的索引:

import numpy as np

data = np.array([[0.9, 0.45, 0.4, 0.35],
[0.4, 0.8, 0.3, 0.25],
[0.5, 0.45, 0.9, 0.35],
[0.2, 0.18, 0.8, 0.1],
[0.6, 0.45, 0.4, 0.9]])

indices = np.array([0,1,2,3,3])

data[np.arange(data.shape[0]), indices]

这输出:

[0.9 0.8 0.9 0.1 0.9]

使用 zip 函数循环遍历两个数组。

def create_array_from_matrix(matrix, indices):
    if len(matrix) != len(indices):
        return None
    res = []
    for row, index in zip(matrix, indices):
        res.append(row[index])

    return res

在第一个数组[0, 1, 2, 3, 3]中,row由每个元素的索引决定,索引处的值为column。这是 enumerate:

的好案例
matrix = [[ ... ], [ ... ], ...] # your matrix
selections = [0, 1, 2, 3, 3]
result = [matrix[i][j] for i, j in enumerate(selections)]

这将比遍历整个 matrix.

更有效