python - 如何将二维数组转为一维数组？

Question

我有二维 numpy 数组 arr:

array([[0., 1., 0.],
       [1., 0., 0.],
       [1., 0., 0.],
       [0., 0., 1.],
       [1., 0., 0.],
       [1., 0., 0.]]

我想将其转换为一维向量:

vec = [1, 0, 0, 2, 0, 0]

值 0、1 和 2 应对应于 art 的列，其中值等于 1。

有没有非for循环的方法来做到这一点？

Answer 1

尝试numpy.nonzero :

>>> import numpy as np
>>> a = np.array([[0., 1., 0.],
...               [1., 0., 0.],
...               [1., 0., 0.],
...               [0., 0., 1.],
...               [1., 0., 0.],
...               [1., 0., 0.]])
>>> np.nonzero(a)[1]
array([1, 0, 0, 2, 0, 0])

或numpy.where :

>>> np.where(a)[1]
array([1, 0, 0, 2, 0, 0])