Оптимизация алгоритма для задачи на перебор

Question

Простая задачка из олимпиады, но мое квадратичное решение не проходит ограничение по времени. Можно ли как нибудь оптимизировать работу моего алгоритма? Условие и свое решение прикладываю ниже :

Лайнландия состоит из n городов, которые расположены на оси Ox. Координата i-го из городов равна xi.

Поликарп прошел от одного города до другого по прямой кратчайшим образом. Он утверждает, что прошел расстояние d.

Выведите количество пар городов, что Поликарп мог пройти от одного к другому. Иными словами, между которыми расстояние составляет ровно d.

Входные данные

В первой строке записаны два целых числа n, d (2 ≤ n ≤ 10^5, 1 ≤ d ≤ 10^9). Вторая строка содержит координаты n городов — последовательность целых чисел x1, x2, ..., xn ( - 10^9 ≤ xi ≤ 10^9). Все xi — различны.

Выходные данные

Выведите искомое количество пар городов. Пары, отличающиеся только порядком городов, следует считать одной и той же парой.

Мое решение:

n,d = map(int, input().split())

l = list(map(int, input().split()))

l.sort()
ans = 0
for i in l:
    if i+d in l:
        ans+=1
print(ans)

Answer 1

Два варианта. Первый с двумя курсорами, второй с ассоциативным массивом.

# x - отсортированный список
def find_pair(x, d, left, right):
    while right < len(x):
        s = x[right] - x[left]
        if s == d:
            return (left, right)
        if s > d:
            left += 1
        if s < d:
            right +=1
    return (-1,-1)

def count1_pairs(x,d):
    x = sorted(x)
    cnt = 0
    l,r = 0,0
    while True:
        l,r = find_pair(x,d,l,r)
        if l < 0:
            break
        else:
            l += 1
            r += 1
            cnt +=1
    return cnt

# Второй вариант
def count2_pairs(x,d):
    mx = { x : i for i,x in enumerate(x)}
    cnt = 0
    for coord in mx.keys():
        if (coord + d) in mx:
            cnt += 1
    return cnt

Тестовые данные: 50 тысяч городов с координатами от 0 до 500 млн. Тестовое расстояние равно расстоянию между городами с номерами 33332 и 10000

np.random.seed(1)
x = np.random.choice(500_000_000, size = 50_000)
d = abs(x[len(x)//3*2] - x[len(x)//5])

Время работы функции count1_pairs : 52.2 ms ± 2.85 ms

Время работы функции count2_pairs : 18.7 ms ± 356 µs

Окончания работы вашей функции я не дождался.

Answer 2

Вот оптимальное по времени решение. Единственное принципиальное отличие от вашего - это замена списка на множество. Поиск в списке требует времени пропорционального длине списка. Поиск в множестве делается за константу.

n,d = map(int, input().split())

l = set(map(int, input().split()))

ans = 0
for i in l:
    if i+d in l:
        ans+=1
print(ans)

P.S. Сортировку я убрал - она не применима к множествам. В вашем варианте сортировка тоже была не нужна.

БЛОГ НА HUSL

Оптимизация алгоритма для задачи на перебор

Ответы (2 шт):