如何使用NumPy提取数组指定列并转换为二进制向量？

实现代码 #

import numpy as np

# 读取数组，若从本地文件读取可替换为你原有的my_data = np.genfromtxt('data_input')
mat = np.array([
    [1., 2., 3., 4., 5.],
    [6., 7., 8., 9., 10.],
    [11., 12., 13., 14., 15.]
])

# 提取第2列（Python索引从0开始，所以取索引1的列）并转为整数
target_col = mat[:, 1].astype(np.int32)
# 定义二进制位长度，示例需要4位
bit_len = 4
# 生成高位到低位的移位权重
shifts = np.arange(bit_len - 1, -1, -1)
# 向量化计算每一位的二进制值
binary_result = ((target_col[:, np.newaxis] >> shifts) & 1).astype(np.int32)

print(binary_result)

运行输出 #

[[0 0 1 0]
 [0 1 1 1]
 [1 1 0 0]]

逻辑说明 #

• 先提取目标列后转换为整数类型，避免浮点型参与位运算报错
• 通过target_col[:, np.newaxis]将一维列向量扩展为二维结构，和移位权重数组广播运算，一次性完成所有元素的各位二进制值计算，不需要循环遍历
• 修改bit_len的数值即可适配不同长度的二进制输出需求

原有代码优化点 #

• 已经单独提取第二列的情况下，无需对整个数组调用np.apply_along_axis遍历所有行，直接处理提取出的单列即可
• 向量化位运算的执行效率远高于自定义函数循环处理，适合大数据量场景

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