Нахождение тех элементов в массиве, которые "близки"

Question

У меня есть 1-мерный отсортированный массив, и я хотел бы найти все пары элементов, разница которых не превышает 5.

Наивным подходом было бы сделать N ^ 2 сравнений, делая что-то вроде

diffs = np.tile(x, (x.size,1) ) - x[:, np.newaxis]
D = np.logical_and(diffs>0, diffs<5)
indicies = np.argwhere(D)

Обратите внимание, что результатом моего примера являются индексы x.Если бы я хотел значения x, которые удовлетворяли бы критериям, я мог бы сделать x[indicies].Это работает для меньших массивов, но не для массивов того размера, с которыми я работаю.

У меня была идея найти места, где между последовательными элементами есть зазоры больше 5.Я бы разделил массив на две части и сравнил все элементы в каждой части.

Является ли это более эффективным способом поиска элементов, которые удовлетворяют моим критериям?Как я мог написать это?

Вот небольшой пример:

x = np.array([ 9, 12, 
           21, 
           36, 39, 44, 46, 47, 
           58, 
           64, 65,])

результат должен выглядеть следующим образом:

array([[ 0,  1],
       [ 3,  4],
       [ 5,  6],
       [ 5,  7],
       [ 6,  7],
       [ 9, 10]], dtype=int64)

Answer 1

Вот решение, которое перебирает смещения при сокращении набора кандидатов до тех пор, пока не останется ни одного:

import numpy as np

def f_pp(A, maxgap):
    d0 = np.diff(A)
    d = d0.copy()
    IDX = []
    k = 1
    idx, = np.where(d <= maxgap)
    vidx = idx[d[idx] > 0]
    while vidx.size:
        IDX.append(vidx[:, None] + (0, k))
        if idx[-1] + k + 1 == A.size:
            idx = idx[:-1]
        d[idx] = d[idx] + d0[idx+k]
        k += 1
        idx = idx[d[idx] <= maxgap]
        vidx = idx[d[idx] > 0]
    return np.concatenate(IDX, axis=0)

data = np.cumsum(np.random.exponential(size=10000)).repeat(np.random.randint(1, 20, (10000,)))

pairs = f_pp(data, 1)
#pairs = set(map(tuple, pairs))

from timeit import timeit
kwds = dict(globals=globals(), number=100)
print(data.size, 'points', pairs.shape[0], 'close pairs')
print('pp', timeit("f_pp(data, 1)", **kwds)*10, 'ms')

Пример выполнения:

99963 points 1020651 close pairs
pp 43.00256529124454 ms

Answer 2

Ваша идея нарезать массив - очень эффективный подход.Поскольку ваши данные отсортированы, вы можете просто вычислить разницу и разделить ее:

d=np.diff(x)
ind=np.where(d>5)[0]
pieces=np.split(x,ind)

Здесь pieces - это список, который вы затем можете использовать в цикле со своим собственным кодом для каждого элемента.

Лучший алгоритм сильно зависит от характера ваших данных, которые я не знаю.Например, другая возможность - написать вложенный цикл:

pairs=[]
for i in range(x.size):
    j=i+1
    while x[j]-x[i]<=5 and j<x.size:
        pairs.append([i,j])
        j+=1

Если вы хотите, чтобы он был более умным, вы можете отредактировать внешний цикл таким образом, чтобы он прыгал, когда j попадает в пробел.