NumPy 高性能技巧：基于多列条件查找最近邻行索引的向量化实现

import numpy as np

# 假设 arr 是从 "data.txt" 加载的原始大型数组
# arr = np.loadtxt("data.txt")

# 示例数据（简化版，模拟原始数据结构）
# 假设原始数据有8列，这里只用4列进行演示
arr = np.array([
    [1.0, 0.0, -1.6, 2.3],  # idx 0
    [1.0, 1.0, -1.6, 2.3],  # idx 1
    [1.0, 0.0, -1.6, 2.3],  # idx 2 (col1=0, col4=2.3, col0=1.0)
    [2.0, 0.0, -1.6, 2.3],  # idx 3
    [1.1, 0.0, -1.6, 2.3],  # idx 4 (col1=0, col4=2.3, col0=1.1)
    [1.2, 0.0, -1.6, 2.3],  # idx 5 (col1=0, col4=2.3, col0=1.2)
    [1.0, 1.0, -1.6, 2.3],  # idx 6 (col1=1, col4=2.3, col0=1.0)
    [1.1, 1.0, -1.6, 2.3],  # idx 7 (col1=1, col4=2.3, col0=1.1)
    [1.0, 0.0, -1.6, 7.3],  # idx 8 (col1=0, col4=7.3, col0=1.0)
    [1.1, 0.0, -1.6, 7.3]   # idx 9 (col1=0, col4=7.3, col0=1.1)
])

# 添加原始索引作为最后一列
# np.c_ 用于按列堆叠数组
arr_indexed = np.c_[arr, np.arange(len(arr))]
print("原始数组（带索引）：\n", arr_indexed)

# 按第二列的值进行排序
arr_sorted_by_col1 = arr_indexed[arr_indexed[:, 1].argsort()]
print("\n按第二列排序后的数组：\n", arr_sorted_by_col1)

# 根据第二列的唯一值进行分块
# np.unique(..., return_index=True) 返回唯一值及其首次出现的索引
unique_col1_values, split_indices = np.unique(arr_sorted_by_col1[:, 1], return_index=True)
# split_indices[1:] 提供了除第一个块之外的所有分界点
splitted_array = np.split(arr_sorted_by_col1, split_indices[1:])

print("\n按第二列分块后的数组（列表形式）：")
for i, block in enumerate(splitted_array):
    print(f"--- 块 {i} (第二列值为 {block[0, 1]}) ---\n", block)

def query_trks(FULL_block, target_col4_value, N):
    """
    在给定的数据块中，查找满足第四列条件且第一列最接近的N个行的原始索引。

    参数:
    FULL_block (np.ndarray): 一个NumPy数组，其中第二列的值是恒定的。
                            最后一列包含原始索引。
    target_col4_value (float): 要匹配的第四列值。
    N (int): 需要查找的最近邻行数。

    返回:
    np.ndarray: 一个数组，包含N个最近邻行的原始索引。
                如果符合条件的行数不足N，则返回所有符合条件的行的索引。
    """
    # 过滤出第四列值与目标值相等的行
    # 注意：这里假设第四列是索引为4，但根据问题描述，它是索引3（0-based）
    # 在示例数据中，我们假设是索引3
    # 根据问题描述，是第四列（索引3）
    filt_mask = (FULL_block[:, 3] == target_col4_value)

    # 获取过滤后的行的索引（在FULL_block内部的相对索引）
    internal_filt_indices = np.where(filt_mask)[0]

    if len(internal_filt_indices) == 0:
        return np.array([], dtype=int) # 没有匹配的行

    # 提取过滤后的行的第一列值
    trks_col0 = FULL_block[internal_filt_indices, 0]

    # 使用广播机制计算所有过滤后的行之间第一列值的绝对差值
    # trks_col0[:, None] 将一维数组转换为列向量，与 trks_col0 进行元素级减法
    # 得到一个 (len(trks_col0), len(trks_col0)) 的差值矩阵
    diff_matrix = np.abs(trks_col0[:, None] - trks_col0)

    # 对差值矩阵按列排序，获取排序后的索引
    # argsort(axis=0) 对每一列进行排序，返回排序后的索引
    # [:N] 获取前N个最小差值的索引
    # 注意：这里需要处理自身与自身的距离为0的情况，通常我们希望排除自身
    # 一种方法是先将对角线元素设为无穷大，或者在获取N个索引后排除自身
    # 考虑到实际需求，通常N个最近邻会包含自身（距离为0），如果需要排除，需额外处理。
    # 这里的argsort会把0距离的自身排在第一位。

    # 为了排除自身，我们可以将对角线元素设置为一个大值
    np.fill_diagonal(diff_matrix, np.inf)

    # 重新排序获取前N个最小差值的索引
    # argsort(axis=0) 会对每一列进行排序，返回排序后的索引
    # [:N] 获取前N个最小差值的索引
    closest_internal_indices = diff_matrix.argsort(axis=0)[:N]

    # 将内部索引转换回 FULL_block 中的实际索引
    # 然后再通过 FULL_block 找到其原始索引

    # 最终结果需要针对每个原始行返回 N 个索引。
    # 这里的 diff_matrix 是针对 trks_col0 的，所以 closest_internal_indices 也是 trks_col0 的索引。
    # 我们需要的是 FULL_block 中对应行的原始索引。

    # 创建一个空的列表来存储每个查询行的N个最近邻原始索引
    result_indices = np.empty((len(internal_filt_indices), N), dtype=int)

    for i, current_row_internal_idx in enumerate(internal_filt_indices):
        # 对于当前行 (FULL_block[current_row_internal_idx])，
        # 找到其在 trks_col0 中的位置 (j)
        # diff_matrix 的第 j 列包含了当前行与其他行的距离

        # 找到当前行在 trks_col0 内部的索引
        j = np.where(internal_filt_indices == current_row_internal_idx)[0][0]

        # 获取第 j 列中 N 个最近邻的内部索引
        n_closest_in_trks_col0 = closest_internal_indices[:, j]

        # 将这些内部索引映射回 FULL_block 的内部索引
        n_closest_in_full_block = internal_filt_indices[n_closest_in_trks_col0]

        # 提取这些行对应的原始索引（最后一列）
        result_indices[i] = FULL_block[n_closest_in_full_block, -1].astype(int)

    return result_indices

# 假设 N = 2
N = 2

# 创建一个列表来存储每个分块计算出的结果
all_computed_indices = []

# 遍历每个分块并调用 query_trks
for block in splitted_array:
    # 在每个block中，第二列的值是相同的，我们只需获取其值
    # 并且对于block中的每一行，我们都需要找到N个最近邻

    # 假设我们为block中的每一行查找最近邻
    # 我们需要迭代block中的每一行作为“当前行”来查询

    # 获取 block 中所有行的原始索引，这些是需要填充结果的行
    original_indices_in_block = block[:, -1].astype(int)

    # 为当前 block 创建一个临时存储结果的数组
    block_results = np.empty((len(block), N), dtype=int)

    for i, row in enumerate(block):
        # 当前行的第二列值和第四列值
        current_col1_val = row[1]
        current_col4_val = row[3]
        current_col0_val = row[0]

        # 过滤出与当前行第四列值相同的行
        # 这里的 FULL_block 实际上就是当前的 block，因为我们已经在外部按 col1 分割了
        # 并且 query_trks 应该只返回与 target_col4_value 匹配的行的最近邻

        # 优化：query_trks 应该接收一个 block 和一个 target_col4_value
        # 并且它应该为 block 中所有满足 target_col4_value 的行计算最近邻

        # 重新设计 query_trks，使其能够处理一个 block 中的所有行
        # 原始的 query_trks 假设 `trk` 是一个值，而不是整个列
        # 让我们回到原始答案的逻辑，它通过 `a[:, 0]` 传入了整个第一列，
        # 这意味着 `query_trks` 内部会为 `trks` 中的每个元素找到最近邻。
        # 这里的 `trks` 是 `FULL[filt, 0]`，即过滤后的第一列。

        # 让我们按照答案的逻辑重新构建 query_trks
        # `query_trks(FULL, trk, v, N)`
        # `FULL` 是当前的 `block`
        # `trk` 在答案中被移除，因为 `trks = FULL[filt, 0]` 已经包含了所有需要比较的第一列值
        # `v` 是 `target_col4_value`