http://stackoverflow.com/questions/8111709/integer-distance

integer-distance.py

#!/usr/bin/env python
from __future__ import print_function
import functools
import sys

try: from collections import Counter
except ImportError:
    pass # ignore

from collections import defaultdict
from itertools import count, islice

try:
    import numpy as np
    import scipy.sparse
    import scipy.spatial
except ImportError: pass # some functionality will be unavailable


def prime_factors_gen():
    """Yield prime factors for each natural number.
Based on
http://stackoverflow.com/questions/567222/simple-prime-generator-in-python/568618#568618

>>> from itertools import islice
>>> list(islice(prime_factors_gen(), 20)) #doctest:+NORMALIZE_WHITESPACE
[(1, []), (2, [2]), (3, [3]), (4, [2]), (5, [5]), (6, [3, 2]),
(7, [7]), (8, [2]), (9, [3]), (10, [5, 2]), (11, [11]), (12, [3, 2]),
(13, [13]), (14, [7, 2]), (15, [5, 3]), (16, [2]), (17, [17]),
(18, [3, 2]), (19, [19]), (20, [5, 2])]
"""
    D = defaultdict(list) # nonprime -> prime factors of `nonprime`
    D[1] = [] # `1` has no prime factors
    for q in count(1): # Sieve of Eratosthenes algorithm
        if q not in D: # `q` is a prime number
            D[q + q] = [q]
            yield q, [q]
        else: # q is a composite
            for p in D[q]: # `p` is a factor of `q`: `q == m*p`
                # therefore `p` is a factor of `p + q == p + m*p` too
                D[p + q].append(p)
            yield q, D[q]
            del D[q]

def multiplicities(n, factors):
    assert n > 0
    for prime in factors:
        alpha = 0 # multiplicity of `prime` in `n`
        q, r = divmod(n, prime)
        while r == 0: # `prime` is a factor of `n`
            n = q
            alpha += 1
            q, r = divmod(n, prime)
        yield prime, alpha

def memoize(func):
    cache = {}
    @functools.wraps(func)
    def wrapper(n):
        try: return cache[n]
        except KeyError:
            ret = cache[n] = func(n)
            return ret
    return wrapper

@memoize
def prime_factors(n):
    return dict(multiplicities(n, uniq_primes_factors[n]))

def distance(a, b):
    """
    a = p1**n1 * p2**n2 * p3**n3 ...
    b = p1**m1 * p2**m2 * p3**m3 ...

    distance = |m1-n1| + |m2-n2| + |m3-n3| ...
    """
    diff = Counter(prime_factors(b))
    diff.subtract(prime_factors(a))
    return sum(abs(d) for d in diff.values())

def init_uniq_primes_factors(limit):
    global uniq_primes_factors
    uniq_primes_factors = dict(islice(prime_factors_gen(), limit))

def get_numbers(size=100):
    numbers = np.random.random_integers(1, 1000000, size=size)
    numbers[0] = 1000000
    return numbers

def solution_gen(numbers):
    for i, ai in enumerate(numbers):
        all_except_i = (aj for j, aj in enumerate(numbers) if j != i)
        yield min(all_except_i, key=lambda x: distance(x, ai))

def numbers2points(numbers):
    """
    n = pK**eK * pL**eL * pM**eM  ->  (...,nK,...,nL,...,nM,...)

    in the base: (p1, p2, ...)
    """
    # primes that are factors of at least one of the numbers
    # mapped to index
    primes2index = dict((p,i) for i, p in enumerate(
        set(prime for n in numbers for prime in uniq_primes_factors[n])))
    M = len(numbers)    # number of points
    N = len(primes2index) # number of coordinates
    points = np.zeros((M,N),dtype=np.int32)
    ##points = scipy.sparse.dok_matrix((M,N), dtype=np.int32)
    for i, n in enumerate(numbers):
        for p, alpha in prime_factors(n).iteritems():
            j = primes2index[p]
            points[i,j] = alpha
    return points

def distance_matrix(points):
    """Calculate distance matrix (i,j) -> distance(a[i],a[j]).

    NOTE: the result is MxM where M is number of points
    """
    return scipy.spatial.distance.squareform(
        scipy.spatial.distance.pdist(points, 'cityblock'))

def solution_gen_distmatrix(numbers):
    points = numbers2points(numbers)
    D = distance_matrix(points)
    D += np.identity(points.shape[0])*D.max() # exclude diagonal (self)
    return (numbers[i] for i in D.argmin(axis=-1))

def solution_gen_kdtree(numbers):
    points = numbers2points(numbers)
    tree = scipy.spatial.cKDTree(points)
    # k=2 self + nearest neighbor
    ind = tree.query(points, k=2, eps=1e-6, p=1)[1]
    return (numbers[i] for i in ind[:,1])


def main():
    N = int(sys.argv[1]) if len(sys.argv) > 1 else 100
    numbers = [1,2,3,4,5,6]
    ##numbers = get_numbers(N)
    # solution_gen():
    # 100000: ~112000 -> 1.5 days
    # 10000:the rest took us 1120.52
    # 1000:the rest took us 11.20
    # 100:the rest took us 0.12
    # 10:the rest took us 0.01

    # solution_gen_distmatrix():
    # 10000: 270 seconds 1GB
    # 1000: 0.4 seconds
    # 100:  0.006 seconds

    # solution_get_kdtree():
    # 10000: 236 1GB
    # 1000: 0.600
    # 100: 0.008
    init_uniq_primes_factors(max(numbers))
    for n in solution_gen(numbers):
        print(n, end=' ')
    print()


##from profilestats import profile
##main = profile(main)
main()