多个二维数组逐坐标求众数的向量化与列表推导实现方法

问题说明 #

现有多个同尺寸二维数组，需统计所有数组对应坐标位置上出现频率最高的值（众数），已知条件：

• 数组元素取值共10-11种
• 单个数组尺寸约为900×900
• 优先需要向量化实现方案，同时确认列表推导式实现的可行性

输入示例 #

A = [[-1, -1, 0, 1, -1], [1, 1, 0, -1, -1], [-1, -1, -1, -1, -1], [-1, 1, -1, -1, 0]]
B = [[-1, -1, 0, 1, -1], [1, -1, 0, -1, -1], [0, 1, -1, 1, -1], [-1, 1, -1, -1, -1]]
C = [[0, -1, 0, 1, -1], [1, -1, 0, -1, -1], [0, 1, -1, 1, -1], [-1, 1, -1, -1, -1]]
D = [[-1, -1, 0, 1, 0], [0, 0, -1, 0, 1], [0, 1, -1, 1, -1], [-1, 1, -1, -1, -1]]

期望输出 #

E = [[-1, -1, 0, 1, -1], [1, -1, 0, -1, -1], [0, 1, -1, 1, -1], [-1, 1, -1, -1, -1]]

向量化实现方案（推荐，适配大尺寸场景） #

基于numpy实现全向量化运算，无Python层显式循环，针对10-11种取值、900×900尺寸的场景，运行耗时仅几十毫秒，性能远高于纯Python循环。

import numpy as np

# 将所有待处理二维数组存入列表，可根据实际情况扩展
arr_list = [A, B, C, D]
# 堆叠为三维数组，维度顺序：(数组数量, 行数, 列数)
stacked_arr = np.array(arr_list, dtype=np.int16)
# 提取所有唯一取值
unique_values = np.unique(stacked_arr)
# 广播统计每个取值在每个坐标位置的出现次数
counts = (stacked_arr == unique_values.reshape(-1, 1, 1)).sum(axis=1)
# 取每个坐标计数最高的取值，转回原生Python二维列表
E = unique_values[np.argmax(counts, axis=0)].tolist()

注：如果遇到同票的情况，该方案会默认返回取值更小的元素，和示例输出逻辑一致。

列表推导式实现方案 #

可以通过嵌套列表推导式实现，但纯Python实现无底层优化，900×900尺寸下运行耗时极长（数分钟级），仅适合小数据量测试使用，不推荐生产场景使用。

# 提前获取数组行列数
rows, cols = len(A), len(A[0])
# 双层列表推导式逐坐标计算众数
E = [
    [
        max(
            # 统计当前坐标所有取值的出现频次
            {v: [arr[i][j] for arr in arr_list].count(v) for v in {arr[i][j] for arr in arr_list}},
            # 取频次最高的值
            key=lambda k: {v: [arr[i][j] for arr in arr_list].count(v) for v in {arr[i][j] for arr in arr_list}}[k]
        )
        for j in range(cols)
    ]
    for i in range(rows)
]

如果要优化列表推导式的性能，可以提前把每个坐标的所有值做聚合缓存，减少重复遍历，但性能依然和numpy向量化方案有量级差距。

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