C closest in array | Python Fiddle

# find C closest elements to an element X in a sorted array A of size S
# bunch of corner cases
#   a) is the element part of the set (diff = 0) ?
#   b) there may not be enough elements ( C > S ) and such corner stuff

from heapq import heapify, heappop, heappush
from itertools import product

from random import randrange

S = randrange(20,100)
C = randrange(1,S/5)
X = randrange(0,S*4)
A = [ randrange(0,S*4) for x in xrange(0,S) ]

print ("%d closest numbers of %d" % (C,X))

# shamelessly stolen from stackoverflow to save time ;-)
# recursive version since it's simpler to understand
# couple bugs corrected in the process ;-)
# extended to deliver the closest element if search fails
def binary_search(value, items, low=0, high=None):
    """
    Binary search function.
    Assumes 'items' is a sorted list.
    The search range is [low, high)
    """

high = len(items) if high is None else high
    pos = low + (high - low) / len(items)

if pos == len(items):
        return pos
    elif items[pos] == value:
        return pos
    elif high == low:
        return high # best approx
    elif items[pos] < value:
        return binary_search(value, items, pos + 1, high)
    else:
        assert items[pos] > value
        return binary_search(value, items, low, pos)

A.sort()

ofs = binary_search(X,A)

# now for the meat of looking for the closest C numbers
# let's walk left and right and pick up the diffs
# it's the simplest version here: run C elems left,
# run C elems right, put the diffs
# on the heap, pull the C abs smallest diffs from the heap
# this is about C * log(C) [ compensated heap, avg C/2 but once insert,
# once pull off ] + 2 * C [loops to put in the heap ]  ~ C  * log(C)

# the minimal O(1) version with alternating left/right running depending
# which side has the smallest element is significantly more code :-}
# but doesn't need the heap of course

diffsheap = []

rofs = ofs
while abs( rofs - ofs ) < C and rofs < len(A):
    # push the abs diff to drive the heap first and sign
    # behind so we actually know what is the absolute number
    heappush(diffsheap,  ( abs ( A[rofs] - X ),
                            + 1 ) )
    rofs += 1
    pass

# careful, right offset already looked @ the found elem itself !
lofs = ofs-1
while abs( ofs - lofs ) < C and lofs >= 0:
    heappush(diffsheap, ( abs( A[lofs] - X ),
                            - 1 ) )
    lofs -= 1
    pass

result = [ heappop(diffsheap) for r in xrange(0,C) ]
# fold over sign of the diff to get absolute diffs
result = [ reduce(lambda x, y: x * y, r) for r in result ]
result.sort()

print ("offsets of best numbers: ", result)

# now let's build heap of diffs of the _whole_ list, then pull off the smallest C elements and see
# whether we got the same

alldiffs = [ abs(A[i] - X) for i in xrange(0,S) ]

alldiffs.sort()

# that's the invariant, we basically make sure all the same offsets in both lists
assert ( sum ( [ abs(r) for r in result ] ) - sum ( alldiffs[:C]) == 0 )

print ("absolute offsets calculated to validate:", alldiffs[:C])

print ("X:", X, "C:", C, "ofs: ", ofs)
print ( "Slice of size C of A around X: ", A[ max(0, ofs-C-1) : min(len(A), ofs+C+1)])
print
print ( "Numbers:", [ X + d for d in result ])

Python Fiddle

Python Cloud IDE