Найти ближайшие k чисел в массиве со сложностью O(logn +k)

Question

Мой код работает, по условию он должен пройти тест с огромными числами менее чем за 10 секунд и сложность должна быть O(logn + k) где k это количество цифр которые нужно возвращать. Подскажите как ускорить мой код, что убрать? Использовать bisect нельзя. Массив сортированный.

import math
def closest(array: list, value: int, count: int):
    masik = list(array)
    start, mid, end, tries = (0, len(masik) // 2, len(masik) - 1, 1)
    array2, final = [], []
    logarray = math.log(len(masik), 2) // 1
    while masik[mid] != value and start <= end:
        if value > masik[mid]:
            start = mid + 1
            tries += 1
        else:
            end = mid - 1
            tries += 1
        mid = (start + end) // 2
        for i in range(1, int(logarray)):
            if tries == logarray - i:
                if len(masik[start:end]) == count or (
                        count <= len(masik[start:end]) <= count ** 2):
                    for i in masik[start + 1:end + 1]:
                        array2.append(i)
        if len(array2) == 0:
            for i in array:
                array2.append(i)
    while len(final) <= count - 1:
        chislo = min(masik, key=lambda x: abs(x - value))
        final.append(chislo)
        masik.remove(chislo)
    final.sort()
    return final

Answer 1

Выполняем бинарный поиск, находя индекс места для вставки нужного значения. O(logn)

Ищем ближайшие значения в массиве справа и слева подобно процедуре слияния из mergesort. O(k)

def getkclosest(a, value, k):
    n = len(a)
    lo = 0
    hi = n
    while lo < hi:
        mid = (lo + hi) // 2
        if a[mid] < value:
            lo = mid + 1
        else:
            hi = mid

    res = []
    hi = lo
    lo -= 1
    cnt = 0

    while cnt < k and lo >= 0 and hi < n:
        if value - a[lo] <= a[hi] - value:
            res.append(a[lo])
            lo -= 1
        else:
            res.append(a[hi])
            hi += 1
        cnt += 1
    while cnt < k and lo >= 0:
        res.append(a[lo])
        lo -= 1
        cnt += 1
    while cnt < k and hi < n:
        res.append(a[hi])
        hi += 1
        cnt += 1
    return res

БЛОГ НА HUSL

Найти ближайшие k чисел в массиве со сложностью O(logn +k)

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