NumPy多行操作与最近邻查找优化指南

有志者，事竟成！如果你在学习文章，那么本文《NumPy优化：多行依赖操作与最近邻查找指南》，就很适合你！文章讲解的知识点主要包括，若是你对本文感兴趣，或者是想搞懂其中某个知识点，就请你继续往下看吧~

本教程详细介绍了如何使用NumPy高效处理复杂的多行依赖操作，以避免性能瓶颈的Python循环。文章核心在于演示如何在一个大型数组中，为每行查找满足特定多列（例如，第二列和第四列值相同）条件的N个最近邻行（基于第一列的数值），并返回其原始索引。通过巧妙地结合数组分割、条件过滤和广播计算，实现了高性能的数据处理。

1. 问题背景与挑战

在数据分析和科学计算中，我们经常需要处理大型数据集，并执行涉及行间依赖的复杂操作。一个典型场景是：给定一个具有多列的大型NumPy数组，我们希望为数组中的每一行找到其他 N 行，这些行满足以下两个条件：

1. 它们在第二列和第四列的值与当前行完全相同。
2. 在满足条件1的所有行中，它们在第一列的值与当前行最为接近。

最终，我们需要一个与原始数组行数相同，列数为 N 的新数组，其中每行包含其对应的 N 个最近邻行的原始索引。

直接使用Python的 for 循环遍历每一行并进行条件判断和查找，对于大型数组来说，会导致显著的性能问题。因此，寻求NumPy的向量化操作是解决此问题的关键。

2. 核心思路与优化策略

为了高效地解决上述问题，我们需要将原始的复杂、全局依赖问题分解为更小、更易于NumPy处理的子问题。核心优化策略包括：

1. 保留原始索引：由于在处理过程中会进行排序和分割，我们需要一种方法来追踪原始行的位置。在数组中添加一列作为原始索引是实现这一点的最佳方式。
2. 按关键列预排序与分割：条件要求第二列和第四列的值相同。我们可以选择其中一列（例如第二列）作为主要分组依据。通过对整个数组按第二列进行排序，然后根据第二列的唯一值将数组分割成多个子数组。这样，每个子数组中的所有行都天然满足了第二列值相同的条件。
3. 在子数组中进行局部处理：对于每个分割后的子数组，我们进一步应用第四列的过滤条件。在满足所有条件的局部数据集中，我们可以高效地使用NumPy的广播机制来计算第一列的距离，并利用 argsort 找到 N 个最近邻的相对索引。
4. 索引转换与结果合并：将局部找到的相对索引转换回原始全局索引，并将所有子数组的结果合并，最终与原始数组关联。

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多行操作与最近邻查找优化指南》的内容就介绍到这里了，是不是学起来一目了然！想要了解更多关于的内容请关注golang学习网公众号！