[英]python, how to select element from each column of matrix
I need to extract one element from each column of a matrix according to an index vector. 我需要根据索引向量从矩阵的每一列中提取一个元素。 Say:
说:
index = [0,1,1]
matrix = [[1,4,7],[2,5,8],[3,6,9]]
Index vector tells me I need the first element from column 1, second element from column 2, and third element from column 3. 索引向量告诉我,我需要第1列的第一个元素,第二列的第2个元素,第3列的第三个元素。
The output should be [1,5,8]
. 输出应为
[1,5,8]
。 How can I write it out without explicit loop? 如何在没有显式循环的情况下将其写出?
Thanks 谢谢
You can use advanced indexing : 您可以使用高级索引 :
index = np.array([0,1,2])
matrix = np.array([[1,4,7],[2,5,8],[3,6,9]])
res = matrix[np.arange(matrix.shape[0]), index]
# array([1, 5, 9])
For your second example, reverse your indices: 对于第二个示例,反转索引:
index = np.array([0,1,1])
matrix = np.array([[1,4,7],[2,5,8],[3,6,9]])
res = matrix[index, np.arange(matrix.shape[1])]
# array([1, 5, 8])
Since you're working with 2-dimensional matrices, I'd suggest using numpy
. 由于您使用的是二维矩阵,因此建议您使用
numpy
。 Then, in your case, you can just use np.diag
: 然后,根据您的情况,您可以使用
np.diag
:
>>> import numpy as np
>>> matrix = np.array([[1,4,7],[2,5,8],[3,6,9]])
>>> matrix
array([[1, 4, 7],
[2, 5, 8],
[3, 6, 9]])
>>> np.diag(matrix)
array([1, 5, 9])
However, @jpp's solution is more generalizable. 但是,@ jpp的解决方案更具通用性。 My solution is useful in your case because you really just want the diagonal of your matrix.
我的解决方案在您的情况下很有用,因为您确实只想要矩阵的对角线。
val = [matrix[i][index[i]] for i in range(0, len(index))]
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.