如何在任意形状numpy ndarray中查找差值最小的两个值及对应索引

实现思路 #

• 先将任意形状的numpy数组展平为一维数组，保留每个元素和原索引的对应关系
• 对展平后的元素按值排序，差值最小的两个元素一定是排序后的相邻元素，无需暴力计算所有两两差值
• 找到差值最小的相邻元素对后，将其对应的一维索引转换为原数组的多维索引即可

代码实现 #

import numpy as np

def find_min_diff_pair_indices(arr: np.ndarray):
    # 展平数组
    flatten_arr = arr.ravel()
    # 获取元素排序后的索引
    sorted_idx = np.argsort(flatten_arr)
    # 计算排序后相邻元素的差值
    adjacent_diffs = np.diff(flatten_arr[sorted_idx])
    # 找到最小差值的位置
    min_diff_pos = np.argmin(adjacent_diffs)
    # 取出最小差值对对应的一维索引
    flat_idx1 = sorted_idx[min_diff_pos]
    flat_idx2 = sorted_idx[min_diff_pos + 1]
    # 转换为原数组的多维索引，转为列表格式返回
    return list(np.unravel_index(flat_idx1, arr.shape)), list(np.unravel_index(flat_idx2, arr.shape))

测试验证 #

# 输入你给出的示例数组
arr = np.array(
[[40.26164428, 63.50590524, 58.30951895],
 [50.99019514, 69.0651866 , 60.44005295],
 [20.24845673, 14.31782106, 58.52349955],
 [54.58937626, 53.03772242, 21.09502311],
 [56.75385449, 57.5847202 ,  1.41421356]])

print(find_min_diff_pair_indices(arr))
# 输出：([0, 2], [2, 2])，和需求完全匹配

补充说明 #

• 方案兼容任意形状的numpy数组，不管是(2,2)、(4,1)还是其他维度都可以正常运行
• 时间复杂度为O(n log n)，n为数组元素总个数，性能远高于暴力两两对比的O(n²)方案
• 如果存在多组相同的最小差值对，把np.argmin替换为np.where(adjacent_diffs == adjacent_diffs.min())遍历处理，即可返回所有符合要求的索引对

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