本教程详细介绍了如何使用numpy高效处理复杂的多行依赖操作，以避免性能瓶颈的python循环。文章核心在于演示如何在一个大型数组中，为每行查找满足特定多列（例如，第二列和第四列值相同）条件的n个最近邻行（基于第一列的数值），并返回其原始索引。通过巧妙地结合数组分割、条件过滤和广播计算，实现了高性能的数据处理。
1. 问题背景与挑战
在数据分析和科学计算中，我们经常需要处理大型数据集，并执行涉及行间依赖的复杂操作。一个典型场景是：给定一个具有多列的大型NumPy数组，我们希望为数组中的每一行找到其他 N 行，这些行满足以下两个条件：
它们在第二列和第四列的值与当前行完全相同。在满足条件1的所有行中，它们在第一列的值与当前行最为接近。
最终，我们需要一个与原始数组行数相同，列数为 N 的新数组，其中每行包含其对应的 N 个最近邻行的原始索引。
直接使用Python的 for 循环遍历每一行并进行条件判断和查找，对于大型数组来说，会导致显著的性能问题。因此，寻求NumPy的向量化操作是解决此问题的关键。
2. 核心思路与优化策略
为了高效地解决上述问题，我们需要将原始的复杂、全局依赖问题分解为更小、更易于NumPy处理的子问题。核心优化策略包括：
保留原始索引：由于在处理过程中会进行排序和分割，我们需要一种方法来追踪原始行的位置。在数组中添加一列作为原始索引是实现这一点的最佳方式。按关键列预排序与分割：条件要求第二列和第四列的值相同。我们可以选择其中一列（例如第二列）作为主要分组依据。通过对整个数组按第二列进行排序，然后根据第二列的唯一值将数组分割成多个子数组。这样，每个子数组中的所有行都天然满足了第二列值相同的条件。在子数组中进行局部处理：对于每个分割后的子数组，我们进一步应用第四列的过滤条件。在满足所有条件的局部数据集中，我们可以高效地使用NumPy的广播机制来计算第一列的距离，并利用 argsort 找到 N 个最近邻的相对索引。索引转换与结果合并：将局部找到的相对索引转换回原始全局索引，并将所有子数组的结果合并，最终与原始数组关联。
3. 逐步实现
接下来，我们将详细介绍如何实现上述策略。
3.1 准备数据
首先，我们需要一个示例数据集。假设我们的数据存储在一个文本文件 data.txt 中，包含多列数据。
import numpy as np# 模拟加载数据，这里使用一个简化的示例数组# 实际应用中，您会从文件加载# data.txt 的格式为：# col0 col1 col2 col3 col4 col5 col6 col7 col8 [desired_output_indices]# 例如：# 1.0 0.0 -1.63e7 0.0 2.39e3 2.39e3 0.0 1.59e-2 0.0 [27, 54]# ...# 假设原始数据是 arr_raw# 为了演示，我们创建一个类似的数据结构# 假设有8列数据arr_raw_data = np.array([ [1.0, 0.0, -1.63711e+07, 0.0, 2.39604e+03, 2.39604e+03, 0.0, 1.59155e-02], # idx 0 [1.0, 1.0, -1.63687e+07, 0.0, 2.39604e+03, 2.39753e+03, 0.0, 1.59280e-02], # idx 1 [1.0, 2.0, -1.63639e+07, 0.0, 2.39604e+03, 2.40051e+03, 0.0, 1.59406e-02], # idx 2 [2.0, 0.0, -1.63711e+07, 0.0, 2.39604e+03, 1.69426e+03, 1.69426e+03, 1.59155e-02], # idx 3 [2.0, 1.0, -1.63694e+07, 1.69426e+03, 2.39604e+03, 1.69575e+03, 1.69426e+03, 1.59280e-02], # idx 4 [6.0, 0.0, -1.63711e+07, 0.0, 7.32003e+03, 7.32003e+03, 0.0, 1.59155e-02], # idx 5 [6.0, 1.0, -1.63638e+07, 0.0, 7.32003e+03, 7.32152e+03, 0.0, 1.59280e-02], # idx 6 [6.0, 2.0, -1.63491e+07, 0.0, 7.32003e+03, 7.32449e+03, 0.0, 1.59406e-02], # idx 7 [1.0, 0.0, -1.63711e+07, 0.0, 2.39604e+03, 2.39604e+03, 0.0, 1.59155e-02], # idx 8 (same as 0) [1.0, 1.0, -1.63687e+07, 0.0, 2.39604e+03, 2.39753e+03, 0.0, 1.59280e-02], # idx 9 (same as 1) [6.0, 0.0, -1.63711e+07, 0.0, 7.32003e+03, 7.32003e+03, 0.0, 1.59155e-02], # idx 10 (same as 5)])# 实际加载数据可能如下：# arr_raw_data = np.loadtxt("data.txt")# 添加原始索引列arr = np.c_[arr_raw_data, np.arange(len(arr_raw_data))]print("原始数组（含索引）：n", arr)
在上述代码中，我们使用 np.c_ 和 np.arange(len(arr_raw_data)) 为数组 arr 添加了一个新的列，其中包含了每一行的原始索引。这一步至关重要，它允许我们在后续的排序和分割操作后，依然能够追溯到原始的行位置。
3.2 预处理：按列排序与分割
为了利用第二列（索引为1）的等值条件，我们首先对数组进行排序，然后根据第二列的唯一值进行分割。
# 按第二列（索引为1）的值进行排序arr = arr[arr[:, 1].argsort()]print("n按第二列排序后的数组：n", arr)# 根据第二列的唯一值分割数组# np.unique(arr[:, 1], return_index=True) 返回唯一值及其首次出现的索引# [1][1:] 用于获取除第一个唯一值以外的所有唯一值的起始索引，作为分割点split_indices = np.unique(arr[:, 1], return_index=True)[1][1:]splitted_array = np.split(arr, split_indices)print("n分割后的子数组列表（基于第二列的唯一值）：")for i, sub_arr in enumerate(splitted_array): print(f"--- 子数组 {i} (第二列值为 {sub_arr[0, 1]}) ---") print(sub_arr)
通过 arr = arr[arr[:, 1].argsort()]，我们将数组按照第二列的值进行升序排列。这使得具有相同第二列值的行彼此相邻。接着，np.split(arr, np.unique(arr[:, 1], return_index=True)[1][1:]) 利用 np.unique 找到第二列所有唯一值的起始索引，并以此作为分割点，将排序后的数组分割成多个子数组。每个子数组内部的所有行都共享相同的第二列值。
3.3 定义 query_trks 函数
query_trks 函数负责在每个子数组内部，根据第四列的条件过滤数据，然后计算第一列的距离并找出 N 个最近邻的原始索引。
def query_trks(FULL, v_val, N): """ 在给定子数组中，根据第四列的特定值过滤， 然后查找N个第一列值最接近的行的原始索引。 参数: FULL (np.ndarray): 当前处理的子数组，已包含原始索引作为最后一列。 v_val (float): 用于过滤第四列的特定值。 N (int): 需要查找的最近邻行数量。 返回: np.ndarray: 包含N个最近邻行的原始索引的数组。 """ # 1. 根据第四列（索引为4）的值进行过滤 # np.where 返回满足条件的索引数组，[0] 获取第一个维度（行索引） filt_indices_in_sub_arr = np.where(FULL[:, 4] == v_val)[0] # 如果过滤后没有足够的行，可以返回空数组或处理 if len(filt_indices_in_sub_arr) < N: # 实际应用中可能需要更复杂的处理，例如填充-1或返回所有可用行 # 这里为了简化，如果不足N个，就返回所有可用行的索引，并用-1填充 if len(filt_indices_in_sub_arr) == 0: return np.full((1, N), -1, dtype=int) # 返回一行N个-1 # 如果有行但不足N个，取所有行，并用-1填充 actual_N = len(filt_indices_in_sub_arr) # 获取第一列的值 trks = FULL[filt_indices_in_sub_arr, 0] # 计算绝对差并排序，找到最近的行 # 这里需要注意，如果 N 比较小，可能需要处理自身的问题 # np.abs(trks[:, None] - trks) 会得到一个 (M, M) 的矩阵 # 如果要排除自身，可以设置对角线为无穷大 diff_matrix = np.abs(trks[:, None] - trks) np.fill_diagonal(diff_matrix, np.inf) # 排除自身 # argsort 找到 N 个最近邻的相对索引 # axis=0 表示对每列（即对每个当前行）进行排序 closest_relative_indices = diff_matrix.argsort(axis=0)[:N] # 将相对索引转换回在FULL中的索引 closest_indices_in_sub_arr = filt_indices_in_sub_arr[closest_relative_indices] # 提取原始索引（最后一列） original_indices = FULL[closest_indices_in_sub_arr, -1] # 填充不足N的部分 result = np.full((original_indices.shape[0], N), -1, dtype=int) result[:, :actual_N] = original_indices[:, :actual_N] return result # 2. 从过滤后的行中提取第一列（索引为0）的值，用于距离计算 trks = FULL[filt_indices_in_sub_arr, 0] # 3. 使用广播计算所有行之间第一列值的绝对差 # trks[:, None] 将 trks 转换为 (M, 1) 形状，trks 仍为 (M,) 形状 # 广播后得到 (M, M) 形状的差值矩阵 diff_matrix = np.abs(trks[:, None] - trks) # 4. 排除自身：将对角线元素设置为无穷大，确保不将当前行作为自己的最近邻 np.fill_diagonal(diff_matrix, np.inf) # 5. 对差值矩阵的每一列（对应每个原始行）进行排序，获取最近邻的索引 # argsort 返回排序后的索引 # [:N] 选取前 N 个最小差值的索引 closest_relative_indices = diff_matrix.argsort(axis=0)[:N] # 6. 将相对索引转换回在 `FULL` 子数组中的实际索引 # filt_indices_in_sub_arr[closest_relative_indices] # 这里的 `closest_relative_indices` 是一个 (N, M) 的数组， # 其中 M 是 `filt_indices_in_sub_arr` 的长度。 # 我们需要将其转置，以便每行对应原始行，每列对应其N个最近邻。 # 并且，这些索引是相对于 `filt_indices_in_sub_arr` 的，需要映射回去。 # 映射回在 FULL 中的索引 mapped_indices_in_full = filt_indices_in_sub_arr[closest_relative_indices] # 7. 提取这些行的原始索引（在 `FULL` 子数组的最后一列） # mapped_indices_in_full 是 (N, M) 形状，其中 M 是子数组中满足条件的行数 # FULL[mapped_indices_in_full, -1] 将提取 (N, M) 形状的原始索引 # 我们需要将其转置为 (M, N) 形状，以便每行对应一个原始行，每列是其一个最近邻 original_indices = FULL[mapped_indices_in_full, -1].T.astype(int) return original_indices
在 query_trks 函数中：
filt_indices_in_sub_arr = np.where(FULL[:, 4] == v_val)[0]：这是在当前子数组中，进一步根据第四列的值进行过滤。v_val 是当前子数组中所有行共享的第四列值。trks = FULL[filt_indices_in_sub_arr, 0]：提取过滤后数据的第一列值，用于计算距离。diff_matrix = np.abs(trks[:, None] – trks)：利用NumPy的广播特性，计算 trks 中所有元素两两之间的绝对差。trks[:, None] 将 trks 转换为列向量，使其能够与行向量 trks 进行广播操作，生成一个差值矩阵。np.fill_diagonal(diff_matrix, np.inf)：将差值矩阵的对角线元素设置为无穷大，目的是排除当前行自身作为其最近邻。closest_relative_indices = diff_matrix.argsort(axis=0)[:N]：argsort(axis=0) 会对 diff_matrix 的每一列（代表原始子数组中的每一行）进行排序，返回排序后的索引。我们取前 N 个索引，即 N 个最近邻的相对索引。original_indices = FULL[mapped_indices_in_full, -1].T.astype(int)：将这些相对索引映射回在 FULL 子数组中的实际行索引，然后从这些行中提取最后一列（原始索引），并进行转置以得到期望的 (M, N) 形状输出。
3.4 整合与结果重组
现在，我们将 query_trks 函数应用到 splitted_array 中的每个子数组，并合并
以上就是NumPy教程：优化多行依赖操作，查找具有共同特征的最近邻行的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1382777.html
微信扫一扫 
支付宝扫一扫 
