[Solved] Find x number of array elements closest to target [closed]

Question

Let’s define some testing data. Let’s have a function:

func getClosest(count: Int, from values: [Int], to targetNumber: Int) -> [Int]

and we will test with the following:

let array = [1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000, 10000, 11000, 12000, 13000, 14000, 15000, 16000, 17000]
print(getClosest(count: 5, from: array, to: 1000))
print(getClosest(count: 5, from: array, to: 10000))

A trivial implementation is to sort the array by distance to the target number and grab the first X values:

func getClosest(count: Int, from values: [Int], to targetNumber: Int) -> [Int] {
    let getDistance: (Int) -> Int = {
        return abs(targetNumber - $0)
    }

    return Array(
        values
        .sorted { getDistance($0) < getDistance($1) }
        .prefix(count)
    )
}

The performance won’t be bad (O(n * log(n)) complexity, the complexity of sorting) but it’s not optimal.

To really get a good complexity, we need to use a specialized data structure called Priority Queue. Priority queue is a name for an abstract data structure that stores elements and can always return the element with the lowest priority.

A naive implementation of such a structure would be:

struct PriorityQueue<Element, Priority: Comparable> {
    private var elements: [(element: Element, priority: Priority)] = []

    var isEmpty: Bool {
        return elements.isEmpty
    }

    var count: Int {
        return elements.count
    }

    mutating func insert(element: Element, priority: Priority) {
        elements.append((element: element, priority: priority))
        elements.sort { $0.priority < $1.priority }
    }

    mutating func peek() -> Element? {
        return elements.first?.element
    }

    mutating func pop() -> Element? {
        guard !isEmpty else {
            return nil
        }

        return elements.remove(at: 0).element
    }

    func getElements() -> [Element] {
        return Array(elements.map { $0.element })
    }
}

This will basically just keep an inner array sorted but note how the interface of a priority queue looks like.

Then we can implement our function this way:

func getClosest(count: Int, from values: [Int], to targetNumber: Int) -> [Int] {
    let getDistance: (Int) -> Int = {
        return abs(targetNumber - $0)
    }

    var queue = PriorityQueue<Int, Int>()

    for value in values {
        queue.insert(element: value, priority: -getDistance(value))
        if queue.count > count {
            _ = queue.pop()
        }
    }

    return queue.getElements()
}

Now, with a naive implementation of priority queue, this won’t help the complexity at all but if we implement the priority queue better then it will work.

A priority queue is usually implemented using Heap, or more specifically Fibonacci heap.

Implementing them is not complicated once you understand the algorithm but there are some implementations already there. See Swift Algorithm Club: Heap and Priority Queue Data Structure for an example.

Accepted Answer