求MatLab ismembertol函数的Julia等效高效实现（百万级数据集）

需求说明 #

我有两个数据集：A（维度N1×3）和B（维度N2×3），需要找出A中那些与B里某几行在各维度独立容差tol=[tol1,tol2,tol3]范围内近似相等的行。

在Matlab中，ismembertol函数可以高效完成这个需求，处理百万级（N1、N2约100万）的数据集仅需数秒，调用代码如下：

ismatch, idx = ismembertol(A,B, 'DataScale', 1, 'ByRow', true, 'OutputAllIndices', true)

但我没找到Julia里的等效高效实现，自己写的朴素双重循环虽然结果正确，但性能极差，代码如下：

# 检查两行是否在给定容差范围内
function row_within_tolerances(row1, row2, tolerances)
    return all(abs.(row1 .- row2) .<= tolerances)
end
# 定义各维度容差
tolerances = [atol1, atol2, atol3]

# 初始化结果数组
tf = falses(size(A, 1))
loc = Vector{Vector{Int}}(undef, size(A, 1))

# 逐行检查A和B的匹配情况
for i in 1:size(A, 1)
    matching_indices = Int[]
    for j in 1:size(B, 1)
        if row_within_tolerances(A[i, :], B[j, :], tolerances)
            push!(matching_indices, j)
        end
    end
    if !isempty(matching_indices)
        tf[i] = true
        loc[i] = matching_indices
    else
        loc[i] = Int[]
    end
end

# 输出结果
println("标记匹配的逻辑数组：")
println(tf)
println("对应A每行在B中的匹配索引数组：")
println(loc)

高效实现方案 #

要处理百万级数据，必须避免O(N1*N2)的双重循环，下面提供两种高效方案：

方案1：使用NearestNeighbors.jl的KD-Tree（通用最优） #

利用KD-Tree进行空间近邻搜索，将各维度容差统一缩放后，搜索距离≤1的邻居即可对应原空间的容差范围：

首先安装依赖包：

using Pkg
Pkg.add("NearestNeighbors")

实现代码：

using NearestNeighbors

function ismembertol_julia(A::Matrix{Float64}, B::Matrix{Float64}, tolerances::Vector{Float64})
    # 按各维度容差缩放数据，统一搜索阈值为1
    B_scaled = B ./ tolerances'
    A_scaled = A ./ tolerances'
    
    # 构建KD-Tree（注意输入需为列向量形式）
    kdtree = KDTree(B_scaled')
    
    # 初始化结果容器
    tf = falses(size(A, 1))
    loc = Vector{Vector{Int}}(undef, size(A, 1))
    
    # 批量搜索每行的匹配项
    for i in 1:size(A_scaled, 1)
        # inrange函数返回所有距离≤阈值的B中行索引
        matching_idxs = inrange(kdtree, A_scaled[i, :], 1.0, true)
        loc[i] = matching_idxs
        tf[i] = !isempty(matching_idxs)
    end
    
    return tf, loc
end

# 调用示例
tolerances = [tol1, tol2, tol3]
tf, loc = ismembertol_julia(A, B, tolerances)

性能说明：KD-Tree构建时间为O(N2 log N2)，单条数据搜索时间为O(log N2)，整体时间复杂度为O(N2 log N2 + N1 log N2)，处理百万级数据的速度与Matlab ismembertol 相当，数秒内即可完成。

方案2：向量化+排序分块（适用于数据有序场景） #

若数据在某维度存在明显有序性，可先对B排序，再通过范围筛选减少比较次数：

function ismembertol_vectorized(A::Matrix{Float64}, B::Matrix{Float64}, tolerances::Vector{Float64})
    # 对B按三维排序，同时保留原索引
    sorted_B = sortslices(hcat(B, 1:size(B,1)), dims=1)
    B_sorted = sorted_B[:, 1:3]
    B_indices = sorted_B[:, 4]
    
    tf = falses(size(A,1))
    loc = Vector{Vector{Int}}(undef, size(A,1))
    
    for i in 1:size(A,1)
        row = A[i,:]
        # 筛选B中各维度在容差范围内的行
        mask = (B_sorted[:,1] .>= row[1]-tolerances[1]) .& (B_sorted[:,1] .<= row[1]+tolerances[1]) .&
               (B_sorted[:,2] .>= row[2]-tolerances[2]) .& (B_sorted[:,2] .<= row[2]+tolerances[2]) .&
               (B_sorted[:,3] .>= row[3]-tolerances[3]) .& (B_sorted[:,3] .<= row[3]+tolerances[3])
        # 获取原索引
        matching_idxs = B_indices[mask]
        loc[i] = matching_idxs
        tf[i] = !isempty(matching_idxs)
    end
    
    return tf, loc
end

性能说明：排序时间为O(N2 log N2)，单条数据通过范围筛选避免全量比较，在数据有序时性能接近KD-Tree方案，但通用性稍弱。

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