Python utility functions for Project Euler

peutils.py

# Most of this code came from
# an external source that has
# long since been forgotten. 
# Feel free to use this
# however you like.
#
# -Nonfrt

import math
import itertools
from decimal import *
from pylab import *
import random
getcontext().prec = 400

################
# Primes Stuff #
################

# Returns a list of primes less than N
def primesBelow(N):
    # http://stackoverflow.com/questions/2068372/fastest-way-to-list-all-primes-below-n-in-python/3035188#3035188
    #""" Input N>=6, Returns a list of primes, 2 <= p < N """
    correction = N % 6 > 1
    N = {0:N, 1:N-1, 2:N+4, 3:N+3, 4:N+2, 5:N+1}[N%6]
    sieve = [True] * (N // 3)
    sieve[0] = False
    for i in range(int(N ** .5) // 3 + 1):
        if sieve[i]:
            k = (3 * i + 1) | 1
            sieve[k*k // 3::2*k] = [False] * ((N//6 - (k*k)//6 - 1)//k + 1)
            sieve[(k*k + 4*k - 2*k*(i%2)) // 3::2*k] = [False] * ((N // 6 - (k*k + 4*k - 2*k*(i%2))//6 - 1) // k + 1)
    return [2, 3] + [(3 * i + 1) | 1 for i in range(1, N//3 - correction) if sieve[i]]

# Returns boolean indicating whether or not n is probably prime
smallprimeset = set(primesBelow(100000))
_smallprimeset = 100000
def isprime(n, precision=4):
    # http://en.wikipedia.org/wiki/Miller-Rabin_primality_test#Algorithm_and_running_time
    if n < 1:
        raise ValueError("Out of bounds, first argument must be > 0")
    elif n <= 3:
        return n >= 2
    elif n % 2 == 0:
        return False
    elif n < _smallprimeset:
        return n in smallprimeset


    d = n - 1
    s = 0
    while d % 2 == 0:
        d //= 2
        s += 1

    for repeat in range(precision):
        a = random.randrange(2, n - 2)
        x = pow(a, d, n)

        if x == 1 or x == n - 1: continue

        for r in range(s - 1):
            x = pow(x, 2, n)
            if x == 1: return False
            if x == n - 1: break
        else: return False

    return True

# https://comeoncodeon.wordpress.com/2010/09/18/pollard-rho-brent-integer-factorization/
def pollard_brent(n):
    if n % 2 == 0: return 2
    if n % 3 == 0: return 3

    y, c, m = random.randint(1, n-1), random.randint(1, n-1), random.randint(1, n-1)
    g, r, q = 1, 1, 1
    while g == 1:
        x = y
        for i in range(r):
            y = (pow(y, 2, n) + c) % n

        k = 0
        while k < r and g==1:
            ys = y
            for i in range(min(m, r-k)):
                y = (pow(y, 2, n) + c) % n
                q = q * abs(x-y) % n
            g = gcd(q, n)
            k += m
        r *= 2
    if g == n:
        while True:
            ys = (pow(ys, 2, n) + c) % n
            g = gcd(abs(x - ys), n)
            if g > 1:
                break

    return g

smallprimes = primesBelow(1000) # might seem low, but 1000*1000 = 1000000, so this will fully factor every composite < 1000000
def primefactors(n, sort=False):
    factors = []

    for checker in smallprimes:
        while n % checker == 0:
            factors.append(checker)
            n //= checker
        if checker > n: break

    if n < 2: return factors

    while n > 1:
        if isprime(n):
            factors.append(n)
            break
        factor = pollard_brent(n) # trial division did not fully factor, switch to pollard-brent
        factors.extend(primefactors(factor)) # recurse to factor the not necessarily prime factor returned by pollard-brent
        n //= factor

    if sort: factors.sort()

    return factors

def factorization(n):
    factors = {}
    for p1 in primefactors(n):
        try:
            factors[p1] += 1
        except KeyError:
            factors[p1] = 1
    return factors

# Returns the number of ints less than n
# that are relatively prime to n
totients = {}
def totient(n):
    if n == 0: return 1

    try: return totients[n]
    except KeyError: pass

    tot = 1
    for p, exp in factorization(n).items():
        tot *= (p - 1)  *  p ** (exp - 1)

    totients[n] = tot
    return tot

def primefactors(n, sort=False):
    factors = []

    limit = int(n ** .5) + 1
    for checker in smallprimes:
        if checker > limit: break
        while n % checker == 0:
            factors.append(checker)
            n //= checker


    if n < 2: return factors
    else : 
        factors.extend(bigfactors(n,sort))
        return factors

def bigfactors(n, sort = False):
    factors = []
    while n > 1:
        if isprime(n):
            factors.append(n)
            break
        factor = pollard_brent(n) 
        factors.extend(bigfactors(factor,sort)) # recurse to factor the not necessarily prime factor returned by pollard-brent
        n //= factor

    if sort: factors.sort()    
    return factors

def divisorGenerator(n):
    factors = factorization(n)
    nfactors = len(factors)
    f = [0] * nfactors
    while True:
        yield reduce(lambda x, y: x*y, [factors.keys()[x]**f[x] for x in range(nfactors)], 1)
        i = 0
        while True:
            f[i] += 1
            if f[i] <= factors[factors.keys()[i]]:
                break
            f[i] = 0
            i += 1
            if i >= nfactors:
                return
    
#####################
# Fibonacci Methods #
#####################

# Helper for calculating the nth Fibonacci
def powLF(n):
    if n == 1:     return (1, 1)
    L, F = powLF(n//2)
    L, F = (L**2 + 5*F**2) >> 1, L*F
    if n & 1:
        return ((L + 5*F)>>1, (L + F) >>1)
    else:
        return (L, F)

# Gives the nth Fibonacci Number
# 1==1th, 1==2nd, 2==3rd, ...
def nthFibi(n):
    if n & 1:
        return powLF(n)[1]
    else:
        L, F = powLF(n // 2)
        return L * F

###################
# Simple Math Ops #
###################

# Returns the greatest common divisor of a and b
def gcd(a, b):
    if a == b: return a
    while b > 0: a, b = b, a % b
    return a

# Returns the least common multiple of a and b
def lcm(a, b):
    return abs((a // gcd(a, b)) * b)

# Returns the sum of x's digits
def digitSum(x):
    Sum = 0
    for i in str(x):
        Sum += int(i)
    return Sum

# The number of combonations of r objects from a set of n
def nCr(n,r):
    if n<r: return 0
    return math.factorial(n)/math.factorial(r)/math.factorial(n-r)

# nth root of val
def nth_root(val, n):
    ret = int(val**(1./n))
    return ret + 1 if (ret + 1) ** n == val else ret

#####################
# Array Opperations #
#####################

# Returns an iterable of all combos of R objects from items
def combosOfLenR(items,R):
    return itertools.combinations(items, R)

# Returns an iterable of all perms of R objects from items
def permsOfLenR(items,R):
    return itertools.permutations(items, R)

# Concatinates to numbers together and returns the result
def concat(x,y):
    return int(str(x)+str(y))

# returns the intersection of 2 lists
def intersection(a,b):
    return [val for val in a if val in b]

# Returns a unique list of the elements in seq, using the idfun
# idfun, if 2 elements have the same id => only one gets into the list
# doesn't work with list of lists
def unique(seq, idfun=None): 
   # order preserving
   if idfun is None:
       def idfun(x): return x
   seen = {}
   result = []
   for item in seq:
       marker = idfun(item)
       # in old Python versions:
       # if seen.has_key(marker)
       # but in new ones:
       if marker in seen: continue
       seen[marker] = 1
       result.append(item)
   return result

# Returns a sorted string version of x
def sort(x):
    return "".join(sorted(str(x)))

# Interchanges two elements, a and b, in list
def switch(a,b,list):
    A, B = list.index(a), list.index(b)
    list[B], list[A] = list[A], list[B]
    return list


#######################
# Continued Fractions #
#######################

# cf helper method
def cfNextStep(decN):
    f = decN.to_integral_exact(rounding=ROUND_FLOOR)
    leftover = decN-Decimal(f)
    return leftover,int(f)

# returns the first 'terms' continuents of the
# square root of n's continued fraction in an array
def cf(n, terms=5, coreOnly=False):
    decN = Decimal(n).sqrt()
    decN,intPart = cfNextStep(decN)
    continuents = [intPart]
    core = []
    for x in xrange(0,terms-1):
        decN,intPart = cfNextStep(Decimal(1)/decN)
        continuents.append(intPart)
        if intPart==2*continuents[0]: # Starts repeating
            core = continuents
            if coreOnly:
                return core
            for x in xrange(0,terms-len(core)-1):
                index = x%(len(core)-1)
                index += 1
                # print index, len(core),n
                continuents.append(core[index])
            break
    return continuents

# returns the kth numerator and demoninator of the continued fraction 
# associated with the Pell's equation derived from D
# [These are possible solns not actual ones]
def pellNumDem(D,k):
    cfSet = cf(D,k)
    num2 = cfSet[0]
    den2 = 1
    if k==1:
        return num2,den2
    num1 = cfSet[1]*cfSet[0]+1
    den1 = cfSet[1]
    if k==2:
        return num1,den1

    num0 = 0
    den0 = 0
    for x in xrange(2,k-1):
        num0 = cfSet[x]*num1+num2
        den0 = cfSet[x]*den1+den2
        num2 = num1
        den2 = den1
        num1 = num0
        den1 = den0
    return num0,den0