NumPy数组分割与搜索操作详解

NumPy 数组分割：`np.array_split()`

NumPy 提供了 np.array_split() 函数，用于将一个一维或多维数组拆分为多个较小的子数组，适合数据划分、交叉验证准备等场景。

语法与核心参数

np.array_split(ary, indices_or_sections, axis=0)

• ary：待分割的输入数组。
• indices_or_sections：定义分割位置。若传入整数 N，则尝试将数组大致均分为 N 份；若传入列表如 [2, 5]，则在对应索引值处切开。
• axis：指定分割的轴向。默认 0（行方向），二维及以上数组可设为 1（列方向）。

基础分割演示

import numpy as np

arr = np.arange(1, 7)  # [1 2 3 4 5 6]

# 按数量均分（元素不足时自动调整）
result = np.array_split(arr, 3)
print(result) 
# 输出: [array([1, 2]), array([3, 4]), array([5, 6])]

# 按指定索引位置切割
result_pos = np.array_split(arr, [2, 4])
print(result_pos) 
# 输出: [array([1, 2]), array([3, 4]), array([5, 6])]

注意：np.array_split() 会尽量保持各子数组长度接近。当数组元素无法整除时，末尾的子数组会容纳剩余元素。该函数始终返回一个 Python 列表，而非连续内存视图。

多维数组与条件切分规范

对于矩阵，可通过设置 axis=1 沿水平方向切分：

mat = np.arange(1, 10).reshape(3, 3)
# mat:
# [[1 2 3]
#  [4 5 6]
#  [7 8 9]]

cols = np.array_split(mat, 3, axis=1)
print(cols)  # 返回三个列切片组成的列表

关于原文提及的“掩码分割”：NumPy 并不直接接受布尔掩码作为 array_split 的参数。若要按条件（如奇偶性）切分数据，标准做法是先通过布尔索引筛选，再分别构建子数组：

data = np.array([1, 2, 3, 4, 5, 6])
even_mask = data % 2 == 0
even_part = data[even_mask]   # [2 4 6]
odd_part  = data[~even_mask]  # [1 3 5]

NumPy 数组搜索：定位与插入索引

条件搜索：np.where()

np.where(condition) 逐元素判断布尔条件，并返回所有满足条件的元素坐标。该函数适用于非结构化数据的快速过滤与索引提取，性能优于原生 Python 循环。

基本检索演示

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

# 查找等于 4 的元素索引
idx_val = np.where(nums == 4)
print(idx_val) 
# 输出: (array([3]),)

# 查找大于 5 的元素索引
idx_gt = np.where(nums > 5)
print(idx_gt)  
# 输出: (array([6, 7, 8]),)

返回值是元组形式，每个维度对应一个索引数组。在多维场景中，元组长度等于数组维度数，可直接用于高级索引 arr[idx_gt[0], idx_gt[1]] 获取子集。

有序表二分搜索：np.searchsorted()

当需要确定某数值在已排序数组中的理论插入位置时，np.searchsorted() 采用二分查找算法，时间复杂度为 O(log n)，在数据量较大时显著优于线性扫描。

语法与侧向控制

np.searchsorted(a, v, side='left')

• a：必须为已升序排列的一维数组。
• v：待查找的标量或兼容类型的数组。
• side：控制相等值时的返回倾向。'left' 返回最左侧插入点（默认），'right' 返回最右侧插入点，常用于维持重复元素的稳定性。

查找演示

sorted_arr = np.array([1, 3, 5, 7, 9])

# 查找 7 的位置
pos_left = np.searchsorted(sorted_arr, 7)
print(pos_left)  # 输出: 3

# 查找 2 的右侧插入点
pos_right = np.searchsorted(sorted_arr, 2, side='right')
print(pos_right)  # 输出: 1

总结

在处理结构化数据时，合理选用分割与搜索函数能显著提升代码执行效率。np.array_split() 负责灵活切分数据空间，np.where() 专注于多条件匹配定位，而 np.searchsorted() 则为有序集合提供极速的二分定位接口。明确三者参数细节与适用边界，可有效覆盖绝大多数基础数据清洗与预处理需求。以上就是NumPy数组分割与搜索函数的详细内容，更多关于NumPy数据处理技巧的资料请关注本站其它相关文章！