如何无循环统计向量/矩阵每行元素右侧更小元素数量？

一、一维向量的处理 #

给定向量如 [3, 1, 5, 7, 9]，要统计每个元素右侧比它小的元素数量，可利用NumPy的广播机制和上三角矩阵筛选实现，完全无需手动写for循环：

import numpy as np

# 输入向量
vec = np.array([3, 1, 5, 7, 9])

# 1. 广播生成两两比较的布尔矩阵：每个元素与所有元素比较是否更大
comp_matrix = vec[:, np.newaxis] > vec

# 2. 提取严格上三角部分（仅保留当前元素右侧的比较结果），每行求和得到统计数
counts = np.sum(np.triu(comp_matrix, k=1), axis=1)

print(counts)  # 输出: [1 0 0 0 0]

原理说明： #

• vec[:, np.newaxis] 将一维向量转为列向量，与原行向量广播后生成 n×n 的布尔矩阵，其中 comp_matrix[i][j] 表示 vec[i] > vec[j]。
• np.triu(comp_matrix, k=1) 保留矩阵的严格上三角区域（k=1 表示排除对角线及下方元素），对应每个元素右侧的所有元素比较结果。
• 对每行求和，即得到该元素右侧比它小的元素总数。

二、二维矩阵的行处理 #

若要对矩阵的每一行执行相同操作，同样可以通过NumPy的广播或内置函数实现，避免手动循环：

方法1：利用 np.apply_along_axis（简洁直观） #

import numpy as np

# 输入矩阵
mat = np.array([[3, 1, 5], [7, 9, 2], [4, 6, 3]])

# 定义单行处理函数
def count_smaller_in_row(row):
    comp = row[:, np.newaxis] > row
    return np.sum(np.triu(comp, k=1), axis=1)

# 对每行应用处理函数
result = np.apply_along_axis(count_smaller_in_row, axis=1, arr=mat)

print(result)
# 输出：
# [[1 0 0]
#  [1 1 0]
#  [1 1 0]]

方法2：三维广播（更高效，适合大矩阵） #

通过扩展维度实现全矩阵的向量化运算，彻底规避循环：

import numpy as np

mat = np.array([[3, 1, 5], [7, 9, 2], [4, 6, 3]])
n_rows, n_cols = mat.shape

# 扩展维度后广播生成三维比较矩阵：(行数, 列数, 列数)
comp_3d = mat[:, :, np.newaxis] > mat[:, np.newaxis, :]

# 提取每行的严格上三角部分并求和
result = np.sum(np.triu(comp_3d, k=1), axis=2)

print(result)  # 输出与方法1一致

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