如何在NumPy中高效实现依赖其他行的行操作（规避循环）

针对你的需求，核心思路是按第2、4列分组，在每个组内对第1列排序后找最近的N个行索引，完全避免逐行循环，大幅提升性能。以下是两种可行实现方式：

方法一：用Pandas（简洁高效，适合大型数组） #

Pandas的groupby是向量化操作，处理分组逻辑非常方便，代码可读性高：

import numpy as np
import pandas as pd

# 生成示例8列数组（替换为你的真实数据）
np.random.seed(42)
arr = np.random.randint(0, 100, size=(10000, 8))
N = 3  # 需要找的最近行数

# 转DataFrame并保留原始索引
df = pd.DataFrame(arr, columns=[f"col{i}" for i in range(8)])
df["original_idx"] = df.index

def get_nearest_in_group(group):
    # 组内按第1列（col0）排序
    sorted_group = group.sort_values("col0").reset_index(drop=True)
    sorted_vals = sorted_group["col0"].values
    sorted_indices = sorted_group["original_idx"].values
    
    # 计算组内元素差值，排除自身（设为无穷大）
    diffs = np.abs(sorted_vals[:, None] - sorted_vals)
    np.fill_diagonal(diffs, np.inf)
    
    # 取每个元素的前N个最小差值对应的索引，映射回原始行号
    nearest_sorted_idx = np.argsort(diffs, axis=1)[:, :N]
    sorted_group["nearest_indices"] = list(sorted_indices[nearest_sorted_idx])
    return sorted_group

# 按第2列（col1）、第4列（col3）分组处理
result_df = df.groupby(["col1", "col3"]).apply(get_nearest_in_group)
# 恢复原始行顺序
result_df = result_df.sort_values("original_idx").reset_index(drop=True)

# 最终结果：每行对应的N个最近行索引
print(result_df["nearest_indices"].values)

方法二：纯NumPy实现（无额外依赖） #

如果不想用Pandas，可通过生成分组键实现向量化分组：

import numpy as np
from scipy.spatial import KDTree  # 可选，提升大组查询效率

arr = np.random.randint(0, 100, size=(10000, 8))
N = 3
nearest_indices = np.empty((arr.shape[0], N), dtype=int)

# 生成分组键：将第2、4列组合为唯一标识（适用于整数类型，浮点数需调整）
keys = arr[:, [1, 3]].view("i8,i8")
unique_keys, inverse = np.unique(keys, return_inverse=True)

# 遍历每个分组（分组数远小于总行数，性能优于逐行循环）
for key_idx in range(len(unique_keys)):
    # 获取当前分组的所有行索引
    group_rows = np.where(inverse == key_idx)[0]
    group_col0 = arr[group_rows, 0].reshape(-1, 1)
    
    # 用KDTree快速找最近邻（比全量差值计算更高效）
    kdtree = KDTree(group_col0)
    # 查询N+1个邻居（第一个是自身，需排除）
    _, idx_in_group = kdtree.query(group_col0, k=N+1)
    nearest_in_group = idx_in_group[:, 1:]
    
    # 将结果写入最终数组
    nearest_indices[group_rows] = group_rows[nearest_in_group]

# 最终结果
print(nearest_indices)

思路可行性验证 #

你之前的问题出在未正确按第2、4列分组，导致过滤的子集不符合条件。现在的思路是完全可行的：

1. 分组确保了每个子集内的行都满足「第2、4列与当前行相同」的条件；
2. 组内排序+最近邻查询，避免了逐行循环，性能随分组数减少呈线性提升；
3. 无论是Pandas还是纯NumPy实现，都是向量化/低循环的操作，适合处理大型数组。

注：如果第2、4列是浮点数，需避免直接用view生成键，可改用np.round取精度后转整数，或用Pandas的groupby直接处理（自动处理浮点数精度问题）。

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