Python中如何使用Numpy将2D数组按位展开构造为3D张量数组

实现代码 #

import numpy as np

# 输入二维正整数矩阵
input_mat = np.array([[1, 2], [3, 4]])

# 自动计算所需二进制位数（由矩阵最大值决定，也可手动指定固定位数）
bit_length = int(input_mat.max()).bit_length()

# 生成3D张量，维度顺序为 [位维度, 原矩阵行, 原矩阵列]，位维度从最高位到最低位排列
output_tensor = ((input_mat[..., np.newaxis] >> np.arange(bit_length-1, -1, -1)) & 1).transpose(2, 0, 1)

print(output_tensor)

输出结果 #

[[[0 0]
  [0 1]]

 [[0 1]
  [1 0]]

 [[1 0]
  [1 0]]]

逻辑说明 #

• 通过bit_length()方法获取矩阵最大值的二进制位数，保证所有元素的二进制位都能被完整展开，可手动替换为固定位数（比如统一展开为8位、16位）
• input_mat[..., np.newaxis]为原二维矩阵新增最后一维，用于后续位维度的广播运算
• 利用右移运算>>和按位与&1提取每个位置对应二进制位的数值
• 最后通过transpose调整维度顺序，将位维度放到第一维，符合示例输出格式
• 若需要最低位排在最前面，只需将移位序列改为np.arange(bit_length)即可

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