如何基于指定行中位数拆分二维NumPy数组？求优雅实现

嘿，刚好做过类似的需求！NumPy并没有专门针对这个场景的内置过滤函数，但用它的矢量化索引特性，就能写出比遍历优雅得多的代码，而且效率还更高（毕竟NumPy的底层是C实现的，比Python循环快太多）。

核心思路 #

1. 先计算指定行的中位数
2. 生成布尔掩码，标记哪些列满足“值≤中位数”，哪些满足“值＞中位数”
3. 直接用布尔掩码对原数组的列进行索引拆分

完整函数实现 #

import numpy as np

def median_split(data, line_number):
    # 计算指定行的中位数
    median_val = np.median(data[line_number])
    # 生成列的布尔掩码
    mask_leq = data[line_number] <= median_val
    mask_gt = data[line_number] > median_val
    # 用掩码索引列，返回两个拆分后的数组
    return data[:, mask_leq], data[:, mask_gt]

举个例子测试 #

比如我们创建一个5行4列的测试数组：

test_data = np.array([
    [1, 2, 3, 4],
    [5, 6, 7, 8],
    [9, 10, 11, 12],
    [13, 14, 15, 16],
    [17, 18, 19, 20]
])
# 按第2行（索引从0开始）拆分
arr_leq, arr_gt = median_split(test_data, 2)
print("≤中位数的列组成的数组：")
print(arr_leq)
print("＞中位数的列组成的数组：")
print(arr_gt)

输出结果会是：

≤中位数的列组成的数组：
[[ 1  2]
 [ 5  6]
 [ 9 10]
 [13 14]
 [17 18]]
＞中位数的列组成的数组：
[[ 3  4]
 [ 7  8]
 [11 12]
 [15 16]
 [19 20]]

为什么这个方法更优？ #

• 完全避免了Python层面的循环遍历，用NumPy的矢量化操作处理，速度快很多，尤其是当数组很大的时候
• 代码简洁直观，一行掩码生成+索引就完成了拆分，可读性拉满

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