如何对同维度NumPy二维数组的匹配列执行逐元素比较？

问题描述 #

我是Python及编程新手，有两个形状相同的NumPy二维数组（均为1000×2000），需要将数组A的每一列与数组B的对应列逐元素比较（即A[:,0]与B[:,0]、A[:,1]与B[:,1]……），而非所有列组合比较。

此前处理一维数组时，我用zip(A,B)配合for循环实现：

import numpy as np

A = np.array([2,5,6,3,7])
B = np.array([1,3,9,4,8])

res_list = []
for number1, number2 in zip(A, B):
    if number1 > number2:
        res_list.append("bigger")
    if number1 < number2:
        res_list.append("smaller")

print(res_list)
# 输出: ['bigger', 'bigger', 'smaller', 'smaller', 'smaller']

但二维数组的实现方式不确定，期望输出是包含2000个子列表的嵌套列表，每个子列表对应一列的比较结果（后续要统计每列中"bigger"和"smaller"的次数）。

输入示例（实际为1000行2000列）：

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

B = np.array([[1, 3, 9, 4, 8],
              [4, 8, 2, 3, 1],
              [3, 7, 1, 8, 9]])

期望输出：

res_list = [["bigger", "bigger", "smaller"], 
            ["bigger", "smaller", "smaller"], 
            ["smaller", "bigger", "bigger"], 
            ["smaller", "smaller", "smaller"], 
            ...]

解决方案 #

方法1：NumPy向量化操作（高效推荐） #

针对1000×2000的大数组，优先用NumPy向量化操作，避免循环大幅提升效率：

import numpy as np

# 生成逐元素比较的掩码数组
mask_bigger = A > B
mask_smaller = A < B

# 创建与原数组形状一致的结果数组，初始为空字符串
result = np.full(A.shape, "", dtype=str)
# 填充对应结果
result[mask_bigger] = "bigger"
result[mask_smaller] = "smaller"

# 转置后转为嵌套列表（原数组行优先，转置后按列分组）
res_list = result.T.tolist()

验证示例输入的输出：

print(res_list)
# 输出:
# [['bigger', 'bigger', 'smaller'],
#  ['bigger', 'smaller', 'smaller'],
#  ['smaller', 'bigger', 'bigger'],
#  ['smaller', 'smaller', 'smaller'],
#  ['smaller', 'bigger', 'smaller']]

方法2：循环实现（新手易理解） #

如果更习惯循环逻辑，可以直接遍历列索引，对对应列逐元素比较：

import numpy as np

res_list = []
num_cols = A.shape[1]  # 获取总列数

for col_idx in range(num_cols):
    col_a = A[:, col_idx]
    col_b = B[:, col_idx]
    sub_list = []
    for a_val, b_val in zip(col_a, col_b):
        if a_val > b_val:
            sub_list.append("bigger")
        elif a_val < b_val:
            sub_list.append("smaller")
        # 若需要处理两值相等的情况，可添加else分支
    res_list.append(sub_list)

统计每列结果次数 #

拿到res_list后，用collections.Counter可以快速统计每列的"bigger"和"smaller"次数：

from collections import Counter

# 生成每列的统计结果列表
column_counts = [Counter(sub_list) for sub_list in res_list]

# 示例：打印第一列的统计结果
print(column_counts[0])
# 输出: Counter({'bigger': 2, 'smaller': 1})

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