numpy.reduce的逆操作：如何将单元素扩展为新维度？

你提到的是reduce的逆操作——将单个元素扩展为多维度数组，Numpy没有直接对应的内置函数，但可以通过向量化的属性提取或数学运算实现，完全规避低效的Python循环。下面针对你给出的两种数组类型分别提供最优解：

一、处理自定义类A的数组 #

对于存储A实例的数组，核心是批量提取元素的value属性，再重塑形状：

实现代码 #

import numpy as np

class A:
    def __init__(self, a, b, c):
        self.value = a, b, c

a = np.array([A(1,2,3), A(4,5,6)])

# 向量化提取属性并重塑
result_a = np.fromiter(
    (x.value for x in a.flat),
    dtype='i,i,i'  # 指定结构化数据类型，对应三个整数
).view(int).reshape(a.shape + (3,))

print(result_a)
# 输出：
# [[1 2 3]
#  [4 5 6]]

原理说明 #

• np.fromiter通过迭代器批量读取元素的value属性，比Python显式循环更高效；
• 用结构化类型'i,i,i'接收元组结果，再通过view(int)转换为普通整数数组；
• reshape(a.shape + (3,))在原数组形状基础上新增最后一个维度，实现扩展。

二、处理打包整数数组 #

对于存储打包整数的数组，直接用向量化位运算提取分量，再堆叠成新维度：

实现代码 #

b = np.array([1<<8 | 2<<4 | 3, 4<<8 | 5<<4 | 6])

# 向量化提取三个分量
component1 = (b >> 8) & 0xF
component2 = (b >> 4) & 0xF
component3 = b & 0xF

# 沿最后一个维度堆叠
result_b = np.stack([component1, component2, component3], axis=-1)

print(result_b)
# 输出：
# [[1 2 3]
#  [4 5 6]]

原理说明 #

• 位运算>>和&都是Numpy支持的向量化操作，直接作用于整个数组，无需循环；
• np.stack将三个一维分量数组沿最后一个维度合并，生成目标形状的二维数组。

三、通用向量化函数 #

整合两种场景，实现自动判断并处理的通用函数：

def expand_vectorized(arr):
    # 检查数组元素类型
    if isinstance(arr.flat[0], A):
        return np.fromiter(
            (x.value for x in arr.flat),
            dtype='i,i,i'
        ).view(int).reshape(arr.shape + (3,))
    else:
        c1 = (arr >> 8) & 0xF
        c2 = (arr >> 4) & 0xF
        c3 = arr & 0xF
        return np.stack([c1, c2, c3], axis=-1)

# 测试两种数组
print(expand_vectorized(a))
print(expand_vectorized(b))

关于np.vectorize的问题 #

你之前用np.vectorize得到数组元组，是因为默认返回类型不匹配。如果一定要用它，需要指定结构化输出类型再转换，但本质上np.vectorize只是Python循环的语法糖，效率远不如上面的向量化方案：

def expanded(element):
    if isinstance(element, A):
        return element.value
    return (element >>8) &0xF, (element>>4)&0xF, element&0xF

# 指定返回结构化类型
f = np.vectorize(expanded, otypes=[('f0', int), ('f1', int), ('f2', int)])
# 转换为普通数组并重塑
result = f(a).view(int).reshape(a.shape + (3,))

内容的提问来源于stack exchange，提问作者Mate de Vita