2
\$\begingroup\$

I'm working on a problem to select k nearest points for a given point. Any advice for bugs, improvements are appreciated, including general advice to implement find nearest k points.

My major idea is using the select-rank algorithm (similar to quicksort, using a pivot value, and divide-conquer the distance array). I think my algorithm time complexity is \$O(n)\$ -- which is \$O(n)\$ to calculate distances, and \$O(\log n)\$ for select rank part. If my calculation is wrong for the time complexity, appreciate for corrections.

import math
import random
def find_nearest_k(distance, start, end, k):
    pivot_index = random.randint(start, end)
    original_start = start
    original_end = end
    while start <= end:
        while start <= end and distance[start][1] <= distance[pivot_index][1]:
            start += 1
        while start <= end and distance[end][1] > distance[pivot_index][1]:
            end -= 1
        if start <= end:
            distance[start], distance[end] = distance[end], distance[start]
        else:
            break
    if start - original_start == k:
        return distance[:start]
    elif start - original_start > k:
        return find_nearest_k(distance, original_start, start - 1, k)
    else:
        return find_nearest_k(distance, start, original_end, k - (start-original_start))

def calculate_distance(center, points):
    result = []
    index = 0
    for point in points:
        result.append((index, (math.pow(center[0]-point[0],2) + math.pow(center[1]-point[1],2))))
        index += 1

    return result

if __name__ == "__main__":
    center = (0,0)
    points = [(1,2),(8,9),(6,5),(2,1),(10,20),(10,9)]
    distances = calculate_distance(center, points)
    print find_nearest_k(distances, 0, len(distances)-1, 2)
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.