如何将shape为(1988,3)的NumPy数组转为对应1所在列索引的1D数组？

可以通过以下步骤实现需求：

1. 先创建一个与原数组行数相同、元素全为NaN的一维数组；
2. 找出原数组中值为1的元素所在的行和列索引；
3. 将对应行的位置赋值为列索引即可。

具体代码实现：

import numpy as np

# 假设preds是你的shape=(1988,3)的numpy数组
preds = np.array([[1, 0, 0],
                  [0, 1, 0],
                  [0, 0, 0],
                  [0, 1, 0],
                  [1, 0, 0],
                  [0, 0, 1]])

# 初始化全为NaN的结果数组
new_preds = np.full(preds.shape[0], np.nan)

# 获取所有值为1的元素的行和列索引
target_rows, target_cols = np.where(preds == 1)

# 给结果数组对应位置赋值
new_preds[target_rows] = target_cols

print(new_preds)
# 输出：[ 0.  1. nan  1.  0.  2.]

补充说明 #

如果原数组中存在某一行有多个1的情况，np.where会返回所有对应的行和列索引，此时new_preds中对应行的最后一个赋值会覆盖之前的。如果需要处理这种场景，可以根据实际需求调整逻辑（比如取第一个出现的1的索引）。

内容的提问来源于stack exchange，提问作者joasa