#! python3 ############################################################################### # # File: enigma.py # RCS: $Header: $ # Description: Useful routines for solving Enigma Puzzles # Author: Jim Randell # Created: Mon Jul 27 14:15:02 2009 # Modified: Tue Feb 21 15:07:32 2023 (Jim Randell) jim.randell@gmail.com # Language: Python (Python 2.7, Python 3.6 - 3.11) # Package: N/A # Status: Free for non-commercial use # URI: http://www.magwag.plus.com/jim/enigma.html # # (c) Copyright 2009-2023, Jim Randell, all rights reserved. # ############################################################################### # -*- mode: Python; python-indent-offset: 2; -*- """ A collection of useful code for solving New Scientist Enigma (and similar) puzzles. The latest version is available at . Currently this module provides the following functions and classes: all_different - check arguments are pairwise distinct all_same - check arguments all have the same value arg - extract an argument from the command line args - extract a list of arguments from the command line as_int - check argument is an integer base2int - convert a string in the specified base to an integer base_digits - get/set digits used in numerical base conversion bit_from_positions - construct an integer by setting bits in specified positions bit_permutations - generate bit permutations bit_positions - return positions of bits set C, nCr - combinatorial function (nCr) cache, cached - decorator for caching functions catch - catch errors in a function call cb - the cube of the argument cbrt - the (real) cube root of a number ceil - generalised ceiling function chain - see: flatten() choose - choose a sequence of values satisfying some functions chunk - go through an iterable in chunks clock - clock arithmetic variant on mod() clump - collect contiguous blocks of the same value collect - collect items according to accept/reject criteria compare - comparator function concat - concatenate a list of values into a string contains - check for contiguous subsequence coprime_pairs - generate coprime pairs cproduct - cartesian product of a sequence of sequences cslice - cumulative slices of an array csum - cumulative sum decompose - construct and call a Decompose() function diff - sequence difference digit_map - create a map of digits to corresponding integer values digrt - the digital root of a number divc - ceiling division divf - floor division div - exact division (or None) divisor - generate the divisors of a number divisor_pairs - generate pairs of divisors of a number divisors - the divisors of a number divisors_pairs - generate pairs of divisors of a number divisors_tuples - generate tuples of divisors of a number dot - vector dot product drop_factors - reduce a number by removing factors dsum - digit sum of a number ediv - exact division (or raise an error) egcd - extended gcd exact_cover - find exact covers from a collection of subsets express - express an amount using specific denominations factor - the prime factorisation of a number factorial - factorial function farey - generate Farey sequences of coprime pairs fcompose - forward functional composition fdiv - float division fib - generate fibonacci sequences filter2 - partition an iterator into values that satisfy a predicate, and those that do not filter_unique - partition an iterator into values that are unique, and those that are not find, rfind - find the index of an object in a sequence find_max - find the maximum value of a function find_min - find the minimum value of a function find_value - find where a function has a specified value find_zero - find where a function is zero first - return items from the start of an iterator flatten - flatten a list of lists flattened - fully flatten a nested structure floor - generalised floor function format_recurring - output the result from recurring() fraction - convert numerator / denominator to lowest terms gcd - greatest common divisor grid_adjacency - adjacency matrix for an n x m grid group - collect values of a sequences into groups hypot - calculate hypotenuse icount - count the number of elements of an iterator that satisfy a predicate implies - logical implication (p -> q) int2base - convert an integer to a string in the specified base int2bcd - convert an integer to binary coded decimal int2roman - convert an integer to a Roman Numeral int2words - convert an integer to equivalent English words intc - ceiling conversion of float to int interleave - interleave values from a bunch of iterators intersect - find the intersection of a collection of containers intf - floor conversion of float to int intr - round a value to the nearest integer invmod - multiplicative inverse of n modulo m ipartitions - partition a sequence with repeated values by index irange - inclusive range iterator irangef - inclusive range iterator with fractional steps iroot - integer kth root function is_coprime - check two numbers are coprime is_cube, is_cube_z - check a number is a perfect cube is_distinct - check a value is distinct from other values is_duplicate - check to see if value (as a string) contains duplicate characters is_pairwise_distinct - check all arguments are distinct is_palindrome - check a sequence is palindromic is_power - check if n = i^k for some integer i is_power_of - check if n = k^i for some integer i is_prime - simple prime test is_prime_mr - Miller-Rabin fast prime test for large numbers is_roman - check a Roman Numeral is valid is_square, is_square_p - check a number is a perfect square is_square_free - check a number is square free is_triangular - check a number is a triangular number isqrt - intf(sqrt(x)) join, joinf - concatenate objects into a string lcm - lowest common multiple M - multichoose function (nMk) map2str - format a map for output match - match a value against a template mgcd - multiple gcd mlcm - multiple lcm mod - return a function to find residues modulo m multiply - the product of numbers in a sequence nconcat - concatenate single digits into an integer ndigits - number of digits used to represent a number in a base nreverse - reverse the digits in an integer nsplit - split an integer into single digits number - create an integer from a string ignoring non-digits ordered - return arguments as an ordered tuple P, nPr - permutations function (nPr) partitions - partition a sequence of distinct values into tuples peek - return an element of a container pi - float approximation to pi poly_* - routines manipulating polynomials, wrapped as Polynomial powers - generate a range of powers prime_factor - generate terms in the prime factorisation of a number prime_factor_rho - generate prime factors of large numbers printf - print with interpolated variables pythagorean_triples - generate Pythagorean triples quadratic - determine roots of a quadratic equation rational - represet a rational number rcompose - reverse functional composition reciprocals - generate reciprocals that sum to a given fraction recurring - decimal representation of fractions recurring2fraction - find the fraction corrresponding to a decimal expansion repdigit - number consisting of repeated digits repeat - repeatedly apply a function to a value reverse - reverse a sequence roman2int - convert a Roman Numeral to an integer rotate - rotate a sequence seq_all_different - check elements of a sequence are pairwise distinct seq_all_same - check elements of a sequence are all the same singleton - return the value from a single valued container split - split a value into characters sprintf - interpolate variables into a string sq - square of x sqrt, sqrtc, sqrtf - the (positive) square root of a number static - decorator to simulate static variables subfactorial - subfactorial function subsets, subseqs - generate subsequences of an iterator substitute - substitute symbols for digits in text substituted_expression - a substituted expression (alphametic/cryptarithm) solver substituted_sum - a solver for substituted sums sum_of_squares - decompose an integer into a sum of squares sumsq - calculate the sum of the squares of a sequence of values tau - tau(n) is the number of divisors of n timed - decorator for timing functions timer - a Timer object translate - substitute values in text tri, T - tri(n) is the nth triangular number trim - remove elements from the start/end of a sequence trirt - the (positive) triangular root of a number tuples - generate overlapping tuples from a sequence ulambda - complex parameter unpacking union - construct the union of a bunch of containers uniq, uniq1 - unique elements of an iterator unzip - inverse of zip unpack - return a function that unpacks its arguments update - create an updated version of a container object zip_eq - check sequences contain the same elements Accumulator - a class for accumulating values CrossFigure - a class for solving cross figure puzzles Decompose - return a decompose() function Delay - a class for the delayed evaluation of a function Denominations - express amounts using specified denominations DominoGrid - a class for solving domino grid puzzles Football - a class for solving football league table puzzles MagicSquare - a class for solving magic squares Matrix - a class for manipulation 2d matrices multiset - a class for manipulating multisets Polynomial - a class for manipulating polynomials Primes - a class for creating prime sieves Rational - select an implementation for rational numbers SubstitutedDivision - a class for solving substituted long division sums SubstitutedExpression - a class for solving general substituted expression (alphametic/cryptarithm) problems SubstitutedSum - a class for solving substituted addition sums Timer - a class for measuring elapsed timings """ # Python 3 style print() and division from __future__ import (print_function, division) __author__ = "Jim Randell " __version__ = "2023-02-25" __credits__ = """Brian Gladman, contributor""" import sys import os import operator import math import functools import itertools import collections import copy import re # maybe use the "six" module for some of this stuff _pythonv = sys.version_info[0:2] # Python version e.g. (2, 7) or (3, 10) if _pythonv[0] == 2: # Python 2.x _python = 2 if _pythonv[1] < 7: print("[enigma.py] WARNING: Python {v} is very old. Things may not work.".format(v=sys.version.split(None, 1)[0])) xrange = xrange reduce = reduce basestring = basestring raw_input = raw_input Sequence = collections.Sequence Iterable = collections.Iterable elif _pythonv[0] > 2: # Python 3.x _python = 3 xrange = range reduce = functools.reduce basestring = str raw_input = input if _pythonv > (3, 6): # Python 3.7 onwards # not: [[ Sequence = collections.abc.Sequence ]] from collections.abc import (Sequence, Iterable) else: Sequence = collections.Sequence Iterable = collections.Iterable # call(, , []) is an alternative to apply() call = lambda fn, args=(), kw=dict(): fn(*args, **kw) # re-exported functions from standard library (to save on imports) defaultdict = collections.defaultdict namedtuple = collections.namedtuple product = itertools.product # cartesian product, but see also: cproduct() # detect if running under PyPy _pypy = getattr(sys, 'pypy_version_info', None) # useful constants enigma = sys.modules[__name__] nl = "\n" pi = math.pi two_pi = pi * 2 inf = float('+inf') empty = frozenset() # the empty set _PY_ENIGMA = os.getenv("PY_ENIGMA") or '' version = lambda: __version__ # add attributes to a function (to use as static variables) # (but for better performance use global variables, or mutable default parameters) def static(**kw): """ simulates static variables in a function by adding attributes to it. static variable in function is accessed as . e.g.: @static(n=0) def gensym(x): gensym.n += 1 return concat(x, gensym.n) >> gensym('foo') 'foo1' >> gensym('bar') 'bar2' >> gensym('baz') 'baz3' (for better performance you can use global variables) """ def _inner(fn): for (k, v) in kw.items(): setattr(fn, k, v) return fn return _inner # useful as a decorator for caching functions (@cached). def cached(f): """ return a cached version of function . the cache can be accessed as attribute 'cache' on function . cache() is also available which will use Python's own function (functools.cache), if available, otherwise cached(). See also: functools.lru_cache() (Python 3.2+), functools.cache() (Python 3.9+). """ f.cache = cache = dict() @functools.wraps(f) def _inner(*k): try: #if k in cache: printf("[{f.__name__}: cache hit, {k}") return cache[k] except KeyError: r = cache[k] = f(*k) #printf("[{f.__name__}: {k} -> {r}]") return r return _inner # or you can use cache, to get functools.cache() (if available) or cached() if not. cache = getattr(functools, 'cache', cached) # wrap a function in another function, e.g. @wrap(uniq, verbose=1) def wrap(fn, *args, **kw): """ a decorator that allows a function to be wrapped in another function. for example: >>> @wrap(uniq) ... def sqmod(n, m): ... for i in irange(1, n): ... yield sq(i) % m will only provide values the first time they are encountered: >>> list(sqmod(10, 10)) [1, 4, 9, 6, 5, 0] """ def _inner(f): @functools.wraps(f) def __inner(*fargs, **fkw): return fn(f(*fargs, **fkw), *args, **kw) return __inner return _inner # the identity function def identity(x): """the identity function: identity(x) == x""" return x # a function that returns a true value def true(*args, **kw): # type: (...) -> bool """a function that ignores any arguments and returns True""" return True # a function that returns its arguments #def tupl(*args, fn=None): return (args if fn is None else fn(args)) # can we treat x as an integer? # include = +/-/0, check for +ve, -ve, 0 def as_int(x, include="", **kw): # type: (...) -> int """ can argument be treated as an integer? can be used to restrict the allowed range, by specifying one or more of: + = allow positive integers - = allow negative integers 0 = allow zero can be specified as a value returned instead of raising an error so things like this work: as_int(0) --> 0 as_int(42) --> 42 as_int(42.0) --> 42 as_int(Fraction(129, 3)) --> 43 as_int(sympy.Integer(42)) --> 42 as_int(sympy.Float(42.0)) --> 42 as_int(sympy.Rational(129, 3)) --> 43 and things like this raise an error: as_int("42") as_int(42.5) as_int(Fraction(129, 2)) as_int(42+0j) as_int(42, include="-") as_int(0, include="+") """ try: n = int(x) if x == n: if include: if n > 0: if '+' in include: return n elif n < 0: if '-' in include: return n else: if '0' in include: return n else: return n return kw['default'] except Exception: pass msg = "invalid integer: " + repr(x) if include: msg += ' [include: ' + include + ']' raise ValueError(msg) # useful routines for solving Enigma puzzles # less than/greater than (or equal) to a target; useful for filter() etc. def lt(t): return (lambda x: x < t) def le(t): return (lambda x: x <= t) def gt(t): return (lambda x: x > t) def ge(t): return (lambda x: x >= t) def between(a, b): return (lambda x: a < x < b) # exclusive between def betweene(a, b): return (lambda x: a <= x <= b) # inclusive between # return a function that increments by a fixed amount def inc(i=1): return (lambda x, i=i: x + i) def mod(m): """ return a function to compute residues modulo . >>> list(map(mod(2), irange(0, 9))) [0, 1, 0, 1, 0, 1, 0, 1, 0, 1] """ return (lambda n: n % m) def clock(m): """ like mod(m) except instead of 0 the function returns . >>> list(map(clock(12), irange(0, 23))) [12, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] """ return (lambda n: (n % m) or m) # like cmp() in Python 2, but results are always -1, 0, +1. # vs can be a triple of values to return instead, default corresponds to (-1, 0, +1) def compare(a, b, vs=None): """ return -1 if a < b, 0 if a == b and +1 if a > b. >>> compare(42, 0) 1 >>> compare(0, 42) -1 >>> compare(42, 42) 0 >>> compare('evil', 'EVIL') 1 """ r = (b < a) - (a < b) return (vs[r + 1] if vs else r) # sign of a number (-1, 0, +1) sign = lambda x: (0 < x) - (x < 0) # = compare(x, 0) # logical implication: p -> q def implies(p, q): """ logical implication: (p -> q) = (not(p) or q) >>> list(p for p in irange(1, 100) if not implies(is_prime(p), p % 6 in (1, 5))) [2, 3] """ return (not p) or q # it's probably quicker (and shorter) to just use: # X not in args # rather than: # is_distinct(X, args) def is_distinct(value, *args): """ check is distinct from values in >>> is_distinct(0, 1, 2, 3) True >>> is_distinct(2, 1, 2, 3) False >>> is_distinct('h', 'e', 'l', 'l', 'o') True """ return value not in args distinct = is_distinct # has distinct values def distinct_values(seq, n=None): seq = list(seq) if n is None: n = len(seq) return len(set(seq)) == n def seq_all_different(seq, fn=None): """ check all elements of are pairwise distinct >>> seq_all_different([0, 1, 2, 3]) True >>> seq_all_different([2, 1, 2, 3]) False >>> seq_all_different(p % 100 for p in primes) False """ seen = set() for x in seq: if fn: x = fn(x) if x in seen: return False seen.add(x) return True # same as distinct_values(args), or distinct_values(args, len(args)) def is_pairwise_distinct(*args, **kw): """ check all arguments are pairwise distinct >>> is_pairwise_distinct(0, 1, 2, 3) True >>> is_pairwise_distinct(2, 1, 2, 3) False >>> is_pairwise_distinct('h', 'e', 'l', 'l', 'o') False """ # this gives the same result as: distinct_values(args, None) # it's probably faster to use a builtin... #return len(set(args)) == len(args) ## even through the following may do fewer tests: #for i in xrange(len(args) - 1): # if args[i] in args[i + 1:]: return False #return True return seq_all_different(args, fn=kw.get('fn', identity)) pairwise_distinct = is_pairwise_distinct all_different = is_pairwise_distinct # returns a Record() with various information on sequence def seq_all_same_r(seq, **kw): """ check to see if a sequence consists of values that are all the same (testing using equality (==)). if a 'value' parameter is passed, elements are checked to see if it is the same as this value, otherwise the elements should be the same as each other. a Record() is returned with the following attributes: same - true if all values in the sequence have the same value value - the value of all elements (or None) empty - if the sequence was empty if the sequence has no elements, a 'value' of None is returned, and 'empty' is set to true. if the sequence contains different values a 'value' of None is returned. if the sequence has fewer than 2 elements, 'same' is trivially true. """ i = iter(seq) n = 0 # is value specified? if 'value' in kw: v = kw['value'] else: # otherwise, use the first value try: v = next(i) n = 1 except StopIteration: # empty sequence return Record(same=True, empty=True, value=None) # check the rest of the sequence for x in i: if x != v: return Record(same=False, empty=False, value=None) n = 1 return Record(same=True, empty=(n == 0), value=v) def seq_all_same(seq, **kw): """ >>> seq_all_same([1, 2, 3]) False >>> seq_all_same([1, 1, 1, 1, 1, 1]) True >>> seq_all_same([1, 1, 1, 1, 1, 1], value=4) False >>> seq_all_same(Primes(expandable=1)) False """ return seq_all_same_r(seq, **kw).same # same as distinct_values(args, 1) def all_same(*args, **kw): """ check all arguments have the same value >>> all_same(1, 2, 3) False >>> all_same(1, 1, 1, 1, 1, 1) True >>> all_same() True """ return seq_all_same_r(args, **kw).same def zip_eq(*ss, **kw): """ check sequences have the same elements. the 'strict' argument is passsed to zip (which supported in some Python versions, and throws an error if the inputs are not of equal length) the 'first' parameter limits checks for the first elements if 'reverse' is set the comparison starts from end >>> zip_eq((1, 2, 3), [1, 2, 3]) True >>> zip_eq((1, 2, 3), [1, 2, 3], irange(1, 3), map(int, "123")) True >>> zip_eq((1, 2, 3, 4), [1, 2, 4, 8]) False >>> zip_eq((1, 2, 3, 4), [1, 2, 4, 8], first=2) True >>> zip_eq((0, 2, 4, 8), [1, 2, 4, 8], reverse=1, first=2) True """ if kw.get('reverse', 0): ss = (reversed(s) for s in ss) z = (zip(*ss, strict=kw['strict']) if 'strict' in kw else zip(*ss)) k = kw.get('first') if k is not None: z = first(z, count=k, fn=iter) for vs in z: if not seq_all_same(vs): return False return True def ordered(*args, **kw): """ return args as a tuple in order. this is useful for making a key for a dictionary. >>> ordered(2, 1, 3) (1, 2, 3) >>> ordered(2, 1, 3, reverse=1) (3, 2, 1) >>> ordered(42) (42,) """ return tuple(sorted(args, **kw)) # I would prefer join() to be a method of sequences: seq.join(sep='') # or a string constructor: str.from_seq(seq, sep='') or just str.join(seq, sep='') # but for now we define a utility function def join(seq, sep='', enc='', fn=str): """ construct a string by joining the items in sequence as strings, separated by separator , and enclosed by the pair . the default separator is the empty string so you can just use: join(seq) instead of: ''.join(seq) >>> join(['a', 'b', 'cd']) 'abcd' >>> join(['a', 'b', 'cd'], sep=',', enc='{}') '{a,b,cd}' >>> join([5, 700, 5]) '57005' """ r = str.join(sep, (fn(x) for x in seq)) if enc: r = enc[0] + r + enc[1] return r def joinf(sep='', enc='', fn=str): "return a joining function" return (lambda x: join(x, sep=sep, enc=enc, fn=fn)) def concat(*args, **kw): """ return a string consisting of the concatenation of the elements of the sequence . the elements will be converted to strings (using str(x)) before concatenation. you can use it instead of str.join() to join non-string lists by specifying a 'sep' argument. >>> concat('h', 'e', 'l', 'l', 'o') 'hello' >>> concat(1, 2, 3, 4, 5) '12345' >>> concat(1, 2, 3, 4, 5, sep=',') '1,2,3,4,5' """ sep = kw.get('sep', '') enc = kw.get('enc', '') if len(args) == 1: try: return join(args[0], sep=sep, enc=enc) except TypeError: pass except: raise return join(args, sep=sep, enc=enc) # reverse a sequence or a map (maybe -> rev()) def reverse(s, fn=None): """ return the reverse of a sequence (str, tuple, list) or map (dict). note: when reversing a map, data may be lost if the original map does not have distinct values. >>> reverse([1, 2, 3]) [3, 2, 1] >>> reverse(first(primes, 6)) [13, 11, 7, 5, 3, 2] >>> reverse("stratagem") 'megatarts' >>> reverse(dict(a=1, b=2, c=3)) {1: 'a', 2: 'b', 3: 'c'} """ ## if it has a 'rev' attribute call that (note: list.reverse() modifies the list) #if hasattr(s, 'rev'): return (s.rev() if fn is None else fn(s.rev())) # if it is a dict, return a reverse map if isinstance(s, dict): return type(s)((v, k) for (k, v) in s.items()) # if it is a string, return a string if fn is None: if isinstance(s, basestring): fn = join elif isinstance(s, tuple): fn = tuple else: fn = list return fn(reversed(s)) rev = reverse # translate text , using map (and optional symbols ) def translate(t, m, s="", embed=1): """ translate the text in according to map (and symbols ) is a string (sequence of letters), if there are sections of the string enclosed in curly braces only those sections will be translated, otherwise the whole string is processed (providing the 'embed' parameter is not disabled). can be: - a dict of -> mappings - a sequence of letters to replace, in which case should give the corresponding substitutions - a function called for each replacement, that should provide the replacement value substitutions can be multiple letters >>> translate("A={A} B={B} C={C}", dict(A=1, B=2, C=3)) 'A=1 B=2 C=3' >>> translate("9567 + 1085 = 10652", "75160892", "DEMNORSY") 'SEND + MORE = MONEY' >>> translate("1->{1}; 2->{2}; 3->{3}", (lambda x: sq(int(x)))) '1->1; 2->4; 3->9' """ t = str(t) # construct the map if callable(m): f = m else: if not isinstance(m, dict): m = dict(zip(m, s)) f = (lambda x: m.get(x, x)) fn = (lambda t: join(map(f, t))) if (not embed) or ('{' not in t): return fn(t) return re.sub(r'{(.*?)}', (lambda x: fn(x.group(1))), t) def nconcat(*digits, **kw): """ return an integer consisting of the concatenation of the list of single digits the digits can be specified as individual arguments, or as a single argument consisting of a sequence of digits. >>> nconcat(1, 2, 3, 4, 5) 12345 >>> nconcat([1, 2, 3, 4, 5]) 12345 >>> nconcat(13, 14, 10, 13, base=16) 57005 >>> nconcat(123,456,789, base=1000) 123456789 >>> nconcat([127, 0, 0, 1], base=256) 2130706433 """ # in Python3 [[ def nconcat(*digits, base=10): ]] is allowed instead base = kw.get('base', 10) # allow a sequence to passed as a single argument if len(digits) == 1 and isinstance(digits[0], (Sequence, Iterable)): digits = digits[0] # this is faster than using: reduce(lambda a, b: a * base + b, digits, 0) n = 0 for d in digits: n *= base n += d return n # or: (slower, and only works with digits < base) #return int(concat(*digits), base=base) # split n into digits, starting with the least significant def nsplitter(n, k=None, base=10): n = abs(as_int(n)) assert base > 1, sprintf("invalid base: {base}") if k is None: # the "natural" number of digits in n while True: (n, r) = divmod(n, base) yield r if n == 0: break else: # the least significant k digits of n for _ in irange(1, k): (n, r) = divmod(n, base) yield r def nsplit(n, k=None, base=10, fn=tuple): """ split an integer into digits (using base representation) if is specified it gives the number of digits to return, if the number has too few digits the the result is zero padded at the beginning, if the number has too many digits then the result includes only the rightmost digits. the sign of the integer is ignored. >>> nsplit(12345) (1, 2, 3, 4, 5) >>> nsplit(57005, base=16) (13, 14, 10, 13) >>> nsplit(123456789, base=1000) (123, 456, 789) >>> nsplit(2130706433, base=256) (127, 0, 0, 1) >>> nsplit(7, 3) (0, 0, 7) >>> nsplit(111**2, 3) (3, 2, 1) """ return fn(reversed(tuple(nsplitter(n, k=k, base=base)))) def dsum(n, k=None, base=10): """ calculate the digit sum of an integer (when represented in the specified base). the sign of the integer is ignored >>> dsum(123456789) 45 >>> dsum(123456789, base=2) 16 """ return sum(nsplitter(n, k=k, base=base)) # population count, Hamming weight, bitsum(), bit_count() if getattr(int, "bit_count", None): dsum2 = int.bit_count else: def dsum2(n): "fast alternative to dsum(n, base=2)"; return bin(abs(n)).count('1', 2) # equivalent to: len(nsplit(n)) # (we could use logarithms for "smallish" numbers) def ndigits(n, base=10): """ return the number of digits in a number, when represented in the specified base. >>> ndigits(factorial(70)) 101 """ #return sum(1 for _ in nsplitter(n, base=base)) return icount(nsplitter(n, base=base)) # maybe -> nrev() def nreverse(n, k=None, base=10): """ reverse an integer (as a digit number using base representation) >>> nreverse(12345) 54321 >>> nreverse(-12345) -54321 >>> nreverse(0xedacaf, base=16) == 0xfacade True >>> nreverse(100) 1 >>> nreverse(1, 3) 100 """ if n < 0: return -nreverse(-n, base=base) else: return nconcat(nsplitter(n, k=k, base=base), base=base) nrev = nreverse from fnmatch import fnmatch # match a value (as a string) to a template # NOTE: match is a soft keyword in Python 3.10+ def match(v, t): """ match a value (as a string) to a template (see fnmatch.fnmatch). to make matching numbers easier if the template starts with a minus sign ('-') then so must the value. if the template starts with a plus sign ('+') then the value must not start with a minus sign. so to match a 4-digit positive number use '+????' as a template. >>> match("abcd", "?b??") True >>> match("abcd", "a*") True >>> match("abcd", "?b?") False >>> match(1234, '+?2??') True >>> match(-1234, '-??3?') True >>> match(-123, '???3') True """ v = str(v) if t.startswith('-'): if v.startswith('-'): (v, t) = (v[1:], t[1:]) else: return False elif t.startswith('+'): if v.startswith('-'): return False else: t = t[1:] return fnmatch(v, t) @static(special={'inf': inf, '+inf': inf, '-inf': -inf}) def number(s, base=10): """ make an integer from a string, ignoring non-digit characters >>> number('123,456,789') 123456789 >>> number('100,000,001') 100000001 >>> number('-1,024') -1024 >>> number('DEAD.BEEF', base=16) == 0xdeadbeef True """ v = number.special.get(s.strip().lower(), None) if v: return v return base2int(s, base=base, strip=1) def split(x, fn=None): """ split into characters (which may be subsequently post-processed by ). >>> split('hello') ['h', 'e', 'l', 'l', 'o'] >>> split(12345) ['1', '2', '3', '4', '5'] >>> list(split(12345, int)) [1, 2, 3, 4, 5] """ return list(map(fn, str(x))) if fn else list(str(x)) # rotate a sequence (move k elements from the beginning to the end) def rotate(s, k=1): """ rotate a sequence by moving elements from the beginning to the end >>> rotate([1, 2, 3, 4], 1) [2, 3, 4, 1] >>> rotate([1, 2, 3, 4], -1) [4, 1, 2, 3] """ return (s if k == 0 else s[k:] + s[:k]) # or you can use itertools.izip_longest(*[iter(l)]*n) for padded chunks def chunk(seq, n=2, pad=0, value=None, fn=tuple): """ iterate through iterable in chunks of size . (for overlapping tuples see tuples()) >>> list(chunk(irange(1, 8))) [(1, 2), (3, 4), (5, 6), (7, 8)] >>> list(chunk(irange(1, 8), 3)) [(1, 2, 3), (4, 5, 6), (7, 8)] """ i = iter(seq) while True: s = fn(itertools.islice(i, 0, n)) if not s: break x = (n - len(s) if pad else 0) yield (s if x == 0 else s + fn([value] * x)) # find contiguous blocks of values (according to fn) def clump(seq, fn=None): """ generate (, ) pairs for contiguous blocks of repeated values in sequence (according to function ). >>> list(clump([1, 1, 1, 2, 2, 3])) [[1, 1, 1], [2, 2], [3]] >>> list(clump("bookkeeper")) [['b'], ['o', 'o'], ['k', 'k'], ['e', 'e'], ['p'], ['e'], ['r']] >>> list(clump(map(tri, irange(1, 10)), fn=mod(2))) [[1, 3], [6, 10], [15, 21], [28, 36], [45, 55]] """ seq = iter(seq) try: v = next(seq) except StopIteration: return xs = [v] if fn is not None: v = fn(v) for x in seq: v_ = (x if fn is None else fn(x)) if v_ == v: xs.append(x) else: yield xs v = v_ xs = [x] yield xs # set union of a bunch of sequences def union(ss, fn=set): """construct a set that is the union of the sequences in """ return fn().union(*ss) # disjoint set union of a bunch of sequences (or None) # any value may appear in only one of the sequences def disjoint_union(ss, fn=set): """ construct a set that is the union of the sequences in . each value in the returned set only appears in one of the sequences (although it may appear multiple times in that sequence). if this is not possible a value of None is returned. >>> disjoint_union([[1], [2], [3], [4]]) == {1, 2, 3, 4} True >>> disjoint_union([[1], [2], [3], [2]]) is None True """ rs = None for xs in ss: if not rs: rs = fn(xs) else: # this seems to be faster than updating rs and checking cardinality for x in xs: if x in rs: return rs.add(x) return rs # set intersection of a bunch of sequences def intersect(ss, fn=set): """construct a set that is the intersection of the sequences in """ i = iter(ss) try: s = fn(next(i)) except StopIteration: pass else: return s.intersection(*i) raise ValueError("empty intersection") # return an element of a container def peek(s, k=0, **kw): """ return an element of a container. empty containers return the specified 'default' value, or raise a ValueError. if k is specified the first k values chosen are discarded (so, for a sequence, you will get s[k]). note that if the container is an iterator, items will be consumed. >>> peek(set([1])) 1 >>> peek([1, 2, 3]) 1 >>> peek("banana") 'b' >>> peek("banana", 4) 'n' >>> peek(primes, 10) 31 >>> peek(p for p in primes if p % 17 == 1) 103 """ if not isinstance(s, dict): # try to index into the container try: return s[k] except (KeyError, IndexError, TypeError): pass # try iterating through the container for (i, x) in enumerate(s): if i == k: return x try: return kw['default'] except KeyError: pass raise ValueError(sprintf("invalid index {k}")) # functions to create a selector for elements/attributes from an object # passing multi=1 forces a multivalued return, even if only one element is specified def item(*ks, **kw): f = operator.itemgetter(*ks) if len(ks) == 1 and kw.get('multi'): return (lambda x: (f(x),)) return f def attr(*ks, **kw): f = operator.attrgetter(*ks) if len(ks) == 1 and kw.get('multi'): return (lambda x: (f(x),)) return f items = lambda n: map(item, xrange(n)) # (x, y, z) = items(3) # select items according to space/comma separated template # item_from("p", "V, L, p") -> item(2) def item_from(select, template, **kw): split = lambda s: (re.split(r'[\s,]+', s.strip("()[]{}")) if isinstance(s, basestring) else s) fields = dict((k, v) for (v, k) in enumerate(split(template))) return item(*(fields[k] for k in split(select)), **kw) def diff(a, b, *rest, **kw): """ return the subsequence of that excludes elements in . >>> diff((1, 2, 3, 4, 5), (3, 5, 2)) (1, 4) >>> join(diff('newhampshire', 'wham')) 'nepsire' """ fn = kw.get('fn', tuple) if rest: b = set(b).union(*rest) return fn(x for x in a if x not in b) # unique combinations: # like uniq(combinations(s, k)) but more efficient def uC(s, k): if k == 0: yield () else: seen = set() for (i, x) in enumerate(s): if x not in seen: for t in uC(s[i + 1:], k - 1): yield (x,) + t seen.add(x) def ucombinations(s, k=None): s = list(s) if k is None: k = len(s) return uC(s, k) # the multiset is implemented as a dict mapping -> class multiset(dict): """ an implementation of multisets. it can be used as an alternative to collections.Counter(), but note the following differences: len() counts the number of elements (not the number of distinct elements) iterating through a multiset provides all elements (not just distinct elements) """ def __init__(self, *vs, **kw): """ create a multiset from one of the following: a dict of -> values a sequence of (, ) values a sequence of individual items (may have repeats) multiple initialisation arguments may be provided, the items from each are added into the multiset. so these are different ways of making the same multiset: multiset("banana") multiset(a=3, b=1, n=2) multiset([('a', 3), ('b', 1), ('n', 2)]) multiset(['b', 'a', 'n', 'a', 'n', 'a']) multiset(dict(a=3, b=1, n=2)) for more control over the initialisation of the multiset you can use: from_dict(), from_pairs(), from_seq() class methods or the corresponding: update_from_dict(), update_from_pairs(), update_from_seq() object methods. """ dict.__init__(self) # deal with any initialisation objects for v in vs: if isinstance(v, dict): # from a dict self.update_from_dict(v) else: # from a sequence v = list(v) try: s = multiset().update_from_pairs(v) self.update_from_dict(s) except (TypeError, ValueError): # maybe more, or maybe just Error self.update_from_seq(v) # add in any keyword items if kw: for (x, n) in kw.items(): self.add(x, n) def update_from_seq(self, vs, count=1): """update a multiset from a sequence of items""" for x in vs: self.add(x, count=count) return self def update_from_pairs(self, vs): """update a multiset from a sequence of (, ) pairs""" for (x, n) in vs: self.add(x, as_int(n)) return self def update_from_dict(self, d): """update a multiset from a dict of -> values""" return self.update_from_pairs(d.items()) @classmethod def from_dict(cls, *vs): """ create a multiset from a dict of -> values (or multiple dicts). """ m = multiset() for v in vs: m.update_from_dict(v) return m @classmethod def from_pairs(self, *vs): """ create a multiset from a sequence of (, ) pairs (or multiple sequences). """ m = multiset() for v in vs: m.update_from_pairs(v) return m @classmethod def from_seq(self, *vs, **kw): """ create a multiset from a sequence of items (or multiple sequences). A keyword argument of 'count' specifies the multiplicity of each element of the sequence inserted into the multiset. """ count = kw.get('count', 1) m = multiset() for v in vs: m.update_from_seq(v, count=count) return m # count all elements in the multiset def size(self): """ the cardinality of the multiset. i.e. a count all the elements in a multiset. to count the number of distinct element types use: s.distinct_size(). this function is used to implement the len() method on multisets. >>> multiset("banana").size() 6 >>> len(multiset("banana")) 6 >>> multiset("banana").distinct_size() 3 """ return sum(dict.values(self)) # len(multiset) == multiset.size() __len__ = size def is_empty(self): return dict.__len__(self) == 0 def is_nonempty(self): return dict.__len__(self) > 0 # is_nonempty is faster than using __len__ __bool__ = __nonzero__ = is_nonempty # return an element def peek(self, k=0, **kw): return peek(self, k=k, **kw) # does this multiset contain elements with multiplicity greater than n? def is_duplicate(self, n=1): return any(v > n for v in dict.values(self)) # does this multiset contain only values with the same multiplicity (n if specified) def all_same(self, n=None): if n is None: return seq_all_same(dict.values(self)) else: return all(v == n for v in dict.values(self)) # all elements of the multiset # (for unique elements use: [[ s.keys() ]]) def elements(self): """ iterate through all elements of the multiset. for distinct elements use: s.keys() this function is used if a multiset is used as an iterator. >>> sorted(multiset("banana")) ['a', 'a', 'a', 'b', 'n', 'n'] >>> sorted(multiset("banana").keys()) ['a', 'b', 'n'] """ for (k, v) in dict.items(self): for _ in xrange(v): yield k __iter__ = elements # the distinct elements # alias for keys() distinct_elements = dict.keys # the number of the distinct elements def __distinct_size(self): """the number of distinct elements in the multiset""" return len(dict.keys(self)) # specialised version for PyPy3 def __distinct_size_pypy3(self): """the number of distinct elements in the multiset""" return len(dict(self)) distinct_size = (__distinct_size_pypy3 if (_pypy and _python > 2) else __distinct_size) # return a count of the item def count(self, item): """return the number of times an item occurs in the multiset""" return dict.get(self, item, 0) # add an item def add(self, item, count=1): """ add an item to a multiset. count can be negative to remove items. """ try: count += self[item] if count == 0: del self[item] return self except KeyError: pass if count < 0: raise ValueError(sprintf("negative count: {item} -> {count}")) if count > 0: self[item] = count return self # remove an item def remove(self, item, count=1): """remove an item from the multiset""" return self.add(item, -count) # like self.items(), but in value order def most_common(self, n=None): """ return the items of the multiset in order of the most common. if n is specifed only the first n items are returned. """ s = sorted(self.items(), key=(lambda t: t[::-1]), reverse=1) return (s if n is None else s[:n]) # provide some useful operations on multisets # update self with some other multisets (item counts are summed) def update(self, *rest): """ update the multiset with some other multisets (or objects that can be interpreted as multisets). item counts are summed. """ for m in rest: if not isinstance(m, dict): m = multiset(m) self.update_from_dict(m) return self # combine self and some other multisets (item counts are summed) def combine(self, *rest): """ return a new multiset that is the result of the original multiset updated with some other multisets (or objects that can be interepreted as multisets). item counts are summed. """ return multiset(self).update(*rest) # union update of self and some other multiset (maximal item counts are retained) def union_update(self, *rest): """ update a multiset with the union of itself and some other multisets (or objects that can be interpreted as multisets). maximal item counts are retained. """ for m in rest: if not isinstance(m, dict): m = multiset(m) for (item, count) in m.items(): self[item] = max(count, self.get(item, 0)) return self # union of self and some other multiset (maximal item counts are retained) def union(self, *rest): """ return a new multiset that is the result of the union of the original multiset and some other multisets (or objects that can be interpreted as multisets). maximal item counts are retained. """ return multiset(self).union_update(*rest) # intersection of self and some other multisets (minimal item counts are retained) def intersection(self, *rest): """ return a new multiset that is the result of the intersection of the original multiset and some other multisets (or objects that can be interpreted as multisets). minimal item counts are retained. """ r = multiset(self) for m in rest: if not isinstance(m, dict): m = multiset(m) r = multiset.from_pairs((item, min(count, r.get(item, 0))) for (item, count) in m.items()) return r # differences between self and m # return (self - m, m - self) def differences(self, m): """ return the differences between self and another multiset m. returns (self - m, m - self) """ if not isinstance(m, dict): m = multiset(m) (d1, d2) = (multiset(), multiset()) for item in set(self.keys()).union(m.keys()): count = self.get(item, 0) - m.get(item, 0) if count > 0: d1.add(item, count) elif count < 0: d2.add(item, -count) return (d1, d2) # difference between self and m # (m may contain items that are not in self, they are ignored) def difference(self, m): """return (self - m)""" return self.differences(m)[0] # is multiset m a subset of self? def issuperset(self, m, strict=0): """test if the multiset contains multiset """ (d1, d2) = self.differences(m) return (not d2) and ((not strict) or bool(d1)) # is multiset m a superset of self? def issubset(self, m, strict=0): """test if the multiset is contained in multiset """ (d1, d2) = self.differences(m) return (not d1) and ((not strict) or bool(d2)) # absolute difference in item counts of the two multisets def symmetric_difference(self, m): """ symmetric difference of this multiset with multiset . the difference in item counts is retained. """ (d1, d2) = self.differences(m) return d1.update(d2) def is_disjoint(self, *rest): """test if the multiset is disjoint from a bunch of other multisets""" for m in rest: if not isinstance(m, dict): m = multiset(m) if any(x in self for x in m): return False return True # multiply item counts def multiply(self, n): """ return a new mutliset derived from the original multiset by multiplying item counts by . """ return multiset.from_pairs((k, n * v) for (k, v) in self.items()) def subsets(self, size=None, min_size=0, max_size=None): """generate subsets of a multiset""" if size is not None: min_size = max_size = size elif max_size is None: max_size = len(self) # distinct elements ks = list(self.keys()) # choose the number of elements for each key for ns in cproduct(irange(0, self[k]) for k in ks): if min_size <= sum(ns) <= max_size: yield multiset.from_pairs(zip(ks, ns)) def copy(self): """return a copy of the multiset""" return multiset.from_dict(self) def min(self, **kw): """ return the minimum item value of a multiset (or ). equivalent to: min(self) """ if (not self) and 'default' in kw: return kw['default'] return min(dict.keys(self)) def max(self, **kw): """ return the maximum item value of a multiset (or ). equivalent to: max(self) """ if (not self) and 'default' in kw: return kw['default'] return max(dict.keys(self)) def sum(self, fn=sum): """ return the sum of items in a multiset. equivalent to: sum(self) """ return fn(v * k for (k, v) in dict.items(self)) def multinomial(self): """ The number of different arrangements of this multiset. >>> multiset("banana").multinomial() 60 """ return multinomial(self.values()) # generate elements in order def sorted(self, key=None, reverse=False): for k in sorted(self.keys(), key=key, reverse=reverse): for _ in xrange(self.get(k)): yield k def map2str(self, sort=1, enc='()', sep=', ', arr='='): """call map2str() on the multiset""" return map2str(self, sort=sort, enc=enc, sep=sep, arr=arr) # generate item pairs def to_pairs(self): return tuple(sorted(self.items())) def to_dict(self): return dict(self) # allow operator overloading on multisets # (let me know if these don't do what you expect) # + = update # - = difference # & = intersection # | = union # * = multiply # < = subset # > = superset __add__ = combine __sub__ = difference __and__ = intersection __or__ = union __iadd__ = update __ior__ = union_update __mul__ = multiply __le__ = issubset __ge__ = issuperset __lt__ = lambda self, m: self.issubset(m, strict=1) __gt__ = lambda self, m: self.issuperset(m, strict=1) def mcombinations(s, k=None): s = sorted(multiset(s)) if k is None: k = len(s) return uC(s, k) # multiset permutations: # a bit like uniq(permutations(s, k)) but more efficient, however items # will not be generated in the same order # # there are more sophisticated algorithms, but this one does the job: # # >>> with Timer(): icount(uniq(subsets("mississippi", select="P"))) # 107899 # [timing] total time: 68.9407666s (68.94s) # # >>> with Timer(): icount(subsets("mississippi", select="mP")) # 107899 # [timing] total time: 0.5661372s (566.14ms) # def mP(d, n, r=()): if n == 0: yield r else: for (k, v) in d.items(): if v > 0: d[k] -= 1 for t in mP(d, n - 1, r + (k,)): yield t d[k] += 1 def mpermutations(s, k=None): s = multiset(s) if k is None: k = len(s) return mP(s, k) # a simple implementation of derangements def derangements(s, k=None): s = list(s) if k is None: k = len(s) for p in itertools.permutations(s, k): if not any(x == s[i] for (i, x) in enumerate(p)): yield p # can be cached() if necessary def subfactorial(n): if n == 0: return 1 return n * subfactorial(n - 1) + (-1 if n % 2 else 1) # subsets (or subseqs) wraps various methods (which can save an import) @static(select_fn=dict(), prepare_fn=dict()) def subsets(s, size=None, min_size=0, max_size=None, select='C', prepare=None, fn=tuple): """ generate tuples representing the subsequences of a (finite) iterator. 'min_size' and 'max_size' can be used to limit the size of the subsequences or 'size' can be specified to produce subsequences of a particular size. the way the elements of the subsequences are selected can be controlled with the 'select' parameter: 'C' = combinations (default) 'P' = permutations 'D' = derangements 'R' = combinations with replacement 'M' = product 'uC' = unique combinations 'mC' = multiset combinations 'mP' = multiset permutations or you can provide your own function. aliases: subseqs(), powerset(). >>> list(subsets((1, 2, 3))) [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)] >>> list(subsets((1, 2, 3), size=2, select='C')) [(1, 2), (1, 3), (2, 3)] >>> list(subsets((1, 2, 3), size=2, select='P')) [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)] >>> list(subsets((1, 2, 3), size=2, select='R')) [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)] >>> list(subsets((1, 2, 3), size=2, select='M')) [(1, 1), (1, 2), (1, 3), (2, 1), (2, 2), (2, 3), (3, 1), (3, 2), (3, 3)] """ if prepare is None: prepare = (list if callable(select) else subsets.prepare_fn.get(select, list)) s = prepare(s) # choose appropriate size parameters if size is not None: if callable(size): size = size(s) min_size = max_size = size elif max_size is None: max_size = len(s) else: if callable(min_size): min_size = min_size(s) if callable(max_size): max_size = max_size(s) # choose an appropriate select function if not callable(select): select = subsets.select_fn[select] # generate the subsets for k in irange(min_size, max_size): for x in select(s, k): yield fn(x) # provide selection functions (where available) # [[ maybe 'R' should be 'M', and 'M' should be 'X' ]] def _subsets_init(): for (k, v, p) in ( ('C', getattr(itertools, 'combinations', None), None), ('P', getattr(itertools, 'permutations', None), None), ('D', derangements, None), ('R', getattr(itertools, 'combinations_with_replacement', None), None), ('M', (lambda s, k: product(s, repeat=k)), None), ('uC', uC, None), ('mC', uC, (lambda s: sorted(multiset(s)))), ('mP', mP, multiset), ): if v: subsets.select_fn[k] = v setattr(subsets, k, v) if p: subsets.prepare_fn[k] = p _subsets_init() # aliases powerset = subsets subseqs = subsets # like filter() but also returns the elements that don't satisfy the predicate # see also partition() recipe from itertools documentation # (but note that itertools.partition() returns (false, true) lists) Filter2 = namedtuple('Filter2', 'true false') def filter2(p, i, fn=list): """ use a predicate to partition an iterable into those elements that satisfy the predicate, and those that do not. returns the partition of the original sequence as: (, ) which can also be accessed from return value r as: r.true r.false alias: partition() >>> tuple(filter2(lambda n: n % 2 == 0, irange(1, 10))) ([2, 4, 6, 8, 10], [1, 3, 5, 7, 9]) """ t = list((x, p(x)) for x in i) return Filter2(fn(x for (x, v) in t if v), fn(x for (x, v) in t if not v)) # alias if you prefer the term partition (but don't confuse it with partitions()) partition = filter2 def is_equal(x, y): """ is_equal(x, y) is the same as (x == y) >>> is_equal(1, 2) False >>> is_equal(42, 42) True >>> is_equal([1, 2, 3], (1, 2, 3)) False """ return (x == y) FilterUnique = namedtuple('FilterUnique', 'unique non_unique') def filter_unique(seq, f=identity, g=identity, st=None): """ for objects in the sequence consider the map f(x) -> g(x) and return a partition of into those objects where f(x) implies a unique value for g(x), and those objects where f(x) implies multiple values for g(x). if the predicate is specified, only objects from the sequence that satisfy the predicate are considered. returns the partition of the original sequence as: (, ) which can also be accessed from return value r as: r.unique r.non_unique See: Enigma 265 alias: partition_unique() "If I told you the first number you could deduce the second" >>> filter_unique([(1, 1), (1, 3), (2, 2), (3, 1), (3, 2), (3, 3)], (lambda v: v[0])).unique [(2, 2)] "If I told you the first number you could not deduce if the second was odd or even" >>> filter_unique([(1, 1), (1, 3), (2, 2), (3, 1), (3, 2), (3, 3)], (lambda v: v[0]), (lambda v: v[1] % 2)).non_unique [(3, 1), (3, 2), (3, 3)] """ # group values by f r = group(seq, st=st, by=f) # collect unique/non-unique items (unq, non) = (list(), list()) if g is identity: # special case if g is not specified for (k, vs) in r.items(): (unq if seq_all_same(vs) else non).extend(vs) else: # general case for (k, vs) in r.items(): (unq if seq_all_same(map(g, vs)) else non).extend(vs) return FilterUnique(unq, non) # alias if you prefer the term partition (but don't confuse it with partitions()) partition_unique = filter_unique def _collect(s, accept, reject, every): for x in s: if (accept is None or accept(x)) and (reject is None or not reject(x)): yield x elif every: raise ValueError(sprintf("collect: failed to collect item: {x}")) def collect(s, accept=None, reject=None, every=0, fn=list): """ collect items from sequence that are accepted by the function (if defined), and not rejected by the function (if defined). return the items that pass the tests (using ) if every=1 then every item must be collected, otherwise None is returned. """ try: return fn(_collect(s, accept, reject, every)) except ValueError: return None def group(seq, by=identity, st=None, f=identity, fn=None): """ group the items of sequence together using the function. items in the same group return the same value when passed to . if the function is specified, only items that satisfy it will be considered. if the function is specified, the function is applied to selected values before they are added to the groups. a dict() is returned where the keys of the dict are the values of the function applied to the items of the sequence, and the values of the dict are the grouped items (collected using , which by default will collect the items in a list, in the order of the original sequence ). >>> group(irange(0, 9), by=mod(2)) {0: [0, 2, 4, 6, 8], 1: [1, 3, 5, 7, 9]} # to reverse a dict into a multivalued map >>> d = dict((n, n % 2) for n in irange(0, 9)) >>> group(d.items(), by=item(1), f=item(0), fn=sorted) {0: [0, 2, 4, 6, 8], 1: [1, 3, 5, 7, 9]} """ d = dict() for x in seq: if st is None or st(x): k = by(x) v = f(x) vs = d.get(k) if vs is None: d[k] = [v] else: vs.append(v) if fn: for (k, vs) in d.items(): d[k] = fn(vs) return d # see ulambda() for a workaround for more complicated unpacking def unpack(fn): """ Turn a function that takes multiple parameters into a function that takes a tuple of those parameters. To some extent this can be used to work around the removal of parameter unpacking in Python 3 (PEP 3113): In Python 2.7 we could write: > fn = lambda (x, y): is_square(x * x + y * y) > fn((3, 4)) 5 but this is not allowed in Python 3. Instead we can write: > fn = unpack(lambda x, y: is_square(x * x + y * y)) > fn((3, 4)) 5 to provide the same functionality. >>> fn = unpack(lambda x, y: is_square(x * x + y * y)) >>> list(filter(fn, [(1, 2), (2, 3), (3, 4), (4, 5)])) [(3, 4)] """ return (lambda args, kw=None: (fn(*args, **kw) if kw else fn(*args))) # unpacked form of zip (which also serves as an inverse to zip) unzip = unpack(zip) # cartesian product of a sequence, cproduct = unpack(itertools.product) def cproduct(ss, **kw): """ the cartesian product of a sequence. so: itertools.product(*()) --> cproduct() itertools.product(As, Bs, Cs) --> cproduct([As, Bs, Cs]) >>> set(cproduct(chunk(irange(1, 4), 2))) == {(1, 3), (1, 4), (2, 3), (2, 4)} True """ return itertools.product(*ss, **kw) # = call(itertools.product, ss, kw) # here's workaround for more complicated parameter unpacking in Python 3 # # in Python 2.7 we could do: # # fn = lambda (x, (y, z)): x + y + z # # instead we can do this: # # fn = ulambda("(x, (y, z))", "x + y + z") # # fn = ulambda("(x, (y, z)): x + y + z") # def ulambda(args, expr=None): """ provide an equivalent to the Python 2 expression: lambda {args}: {expr} in Python 3 where {args} specifies a complex parameter unpacking of arguments e.g.: >>> dist = ulambda("(x1, y1), (x2, y2): hypot(x2 - x1, y2 - y1)") >>> dist((1, 2), (5, 5)) 5.0 """ if expr is None: (args, _, expr) = (x.strip() for x in args.partition(":")) if _python == 2: # in Python 2 it is straightforward lambda expr = sprintf("lambda {args}: {expr}") else: # in Python 3 any of the following achieve the same #expr = sprintf("lambda *_x_: [{expr} for [{args}] in [_x_]][0]") #expr = sprintf("lambda *_x_: peek({expr} for [{args}] in [_x_])") expr = sprintf("lambda *_x_: next({expr} for [{args}] in [_x_])") return eval(expr) # count the number of occurrences of a predicate in an iterator def icount(i, p=None, t=None): """ count the number of elements in iterator that satisfy predicate

, the termination limit controls how much of the iterator we visit, so we don't have to count all occurrences. So, to find if exactly elements of satisfy

use: icount(i, p, n + 1) == n which is what icount_exactly(i, p, n) does. This will examine all elements of to verify there are exactly 4 primes less than 10: >>> icount_exactly(irange(1, 10), is_prime, 4) True But this will stop after testing 73 (the 21st prime): >>> icount_exactly(irange(1, 100), is_prime, 20) False To find if at least elements of satisfy

use: icount(i, p, n) == n This is what icount_at_least(i, p, n) does. The following will stop testing at 71 (the 20th prime): >>> icount_at_least(irange(1, 100), is_prime, 20) True To find if at most elements of satisfy

using Kroneker's method def poly_factor(p, F=None, div=None): if F is None: F = Rational() if div is None: div = Rdiv(F) # first find factors of x (n, p) = (0, list(p)) while p[0] == 0: n += 1 p.pop(0) if n > 0: yield ([0, 1], n) # find other linear factors for x in poly_rational_roots(p, domain='Q', F=F): f = [-int(x.numerator), int(x.denominator)] (n, p) = poly_div(p, f, div=div) if n > 0: yield (f, n) # look for factors of degree k k = 2 while len(p) > 2 * k: q = poly_scale(p) # evaluate q at (k + 1) values, and record the divisors ds = list() for i in irange(0, k): vs = divisors(poly_value(q, i)) if i > 0: vs += list(-x for x in vs) ds.append(vs) # choose potential values for polynomial factor at the values for vs in cproduct(ds): if mgcd(*vs) != 1: continue # interpolate the polynomial factor f = poly_interpolate(enumerate(vs), field=F) if len(f) - 1 < k: continue (n, p) = poly_div(p, f, div=div) if n > 0: yield (f, n) break else: k += 1 # anything left? if len(p) > 1: f = poly_scale(p) yield (f, 1) p = [div(p[-1], f[-1])] if p != poly_unit: yield (p, 1) # drop factors in qs from polynomial p def poly_drop_factor(p, *qs): for q in qs: while True: (d, r) = poly_divmod(p, q) if r != poly_zero: break p = d return p # calculate cyclotomic polynomials @static(cache={1: [-1, 1]}, cache_enabled=1) def poly_cyclotomic(n, fs=None, div=rdiv, fn=prime_factor): """ return the nth cyclotomic polynomial >>> poly_cyclotomic(7) [1, 1, 1, 1, 1, 1, 1] >>> poly_cyclotomic(12) [1, 0, -1, 0, 1] >>> poly_cyclotomic(30) [1, 1, 0, -1, -1, -1, 0, 1, 1] """ if n < 1: return None r = poly_cyclotomic.cache.get(n) if r is None: if fs is None: fs = list(fn(n)) if len(fs) == 1: (p, e) = fs[0] if e == 1: # n is prime r = [1] * n else: # power of a prime q = n // p r = poly_from_pairs((k * q, 1) for k in irange(0, p - 1)) elif fs[0] == (2, 1): # 2n, invert the odd positions in cyclotomic[n] r = list(poly_cyclotomic(n // 2, fs=fs[1:], div=div, fn=fn)) for i in range(1, len(r), 2): r[i] = -r[i] else: # C[n] = multiply((x^d - 1) ^ mobius(n // d) for d in divisors(n)) # we can specialise multiplication and division by (x^d - 1) (r, ds) = ([1], []) for d in multiples(fs): m = mobius(n // d, fn=fn) if m == 1: # r *= (x^d - 1) r = poly_sub([0] * d + r, r) elif m == -1: ds.append(d) for d in ds: # r /= (x^d - 1) r_ = list() for c in reversed(r[d:]): r_.insert(0, c) if len(r_) > d: r_[0] += r_[d] r = r_ if poly_cyclotomic.cache_enabled: poly_cyclotomic.cache[n] = r return r # wrap the whole lot up in a class class Polynomial(list): """ A class for manipulating polynomials in one variable. Polynomials are represented by a list of their coefficents: a + b.x + c.x^2 + d.x^3 + ... -> [a, b, c, d, ...] """ def __repr__(self, var='x'): return self.__class__.__name__ + "[" + poly_print(self, var=var) + "]" def print(self, var='x'): return poly_print(self, var=var) def __hash__(self): return hash(tuple(self)) def __add__(self, other): if not isinstance(other, Polynomial): other = Polynomial([other]) return self.__class__(poly_add(self, other)) # this allows: + (e.g. 3 + p) __radd__ = __add__ def __iadd__(self, other): return self + other def __mul__(self, other): if isinstance(other, Polynomial): # multiply polynomials return self.__class__(poly_mul(self, other)) else: # multiply coefficients return self.__class__(other * c for c in self) # this allows: * (e.g. 3 * p) __rmul__ = __mul__ def __neg__(self): return self.__class__(poly_neg(self)) def __pow__(self, n): return self.__class__(poly_pow(self, n)) def __sub__(self, other): if not isinstance(other, Polynomial): other = Polynomial([other]) return self.__class__(poly_sub(self, other)) # this allows: - (e.g. 3 - p) __rsub__ = lambda self, other: -self + other __call__ = poly_value def copy(self): "return a copy of the polynomial" return self.__class__(self) def degree(self): "return the degree of the polynomial" return len(self) - 1 def coeff(self, k, default=0): "return the coefficient of x^k in the polynomial" if 0 <= k < len(self): return self[k] else: return default def to_pairs(self): "an iterator that returns (, ) pairs of the polynomial" for p in enumerate(self): if p[1] != 0: yield p def is_zero(self): "check if this polynomial is zero: p(x) = 0" return self == poly_zero def is_unit(self): "check if this polynomial is the unit polynomial: p(x) = 1" return self == poly_unit def map(self, fn): """ return a polynomial that is the result of applying to the coefficents in this polynomial """ return self.__class__(poly_map(self, fn)) def scale(self): """ return a polynomial with integer coefficients derived from this polynomial by multiplying the coefficients by a scalar value """ return self.__class__(poly_scale(self)) def derivative(self, k=1): "return the derivative of the polynomial" return self.__class__(poly_derivative(self, k)) def integral(self, c=0, div=rdiv): "return the indefinite integral of the polynomial" return self.__class__(poly_integral(self, c=c, div=div)) def rational_roots(self, v=0, domain="Q", include="+-0", F=None): "find rational roots of the polynomial" p = (self - v if v else self) return poly_rational_roots(p, domain=domain, include=include, F=F) def divmod(self, q, div=rdiv): (d, r) = poly_divmod(self, q, div=div) return (self.__class__(d), self.__class__(r)) def compose(self, other): "return a polynomial which is the result of the applying this polynomial to another" return self.__class__(poly_compose(self, other)) def factor(self, F=None, div=None): "generate factors of the polynomial" for (f, n) in poly_factor(self, F=F, div=div): yield (self.__class__(f), n) def drop_factor(self, *qs): return self.__class__(poly_drop_factor(self, *qs)) @classmethod def from_pairs(cls, ps): return cls(poly_from_pairs(ps)) @classmethod def unit(cls): return cls(poly_unit) @classmethod def zero(cls): return cls(poly_zero) # sum() is only documented for "numeric" values (although it works) # but you can use this instead... @classmethod def sum(cls, ps): "return the sum of a sequence of polynomials" return cls(poly_sum(ps)) @classmethod def multiply(cls, ps): "return the product of a sequence of polynomials" return cls(poly_multiply(ps)) @classmethod def interpolate(cls, ps, field=None): "return a polynomial that fits the (x, y) values in " r = poly_interpolate(ps, field=field) return (None if r is None else cls(r)) @classmethod def cyclotomic(cls, n, div=rdiv, fn=prime_factor): "return the nth cyclotomic polynomial" p = poly_cyclotomic(n, div=div, fn=fn) return (None if p is None else cls(p)) ############################################################################### # Prime Sieves _primes_array = bytearray _primes_size = 1295 _primes_chunk = lambda n: 2 * n _primes_array_optimise = 1 # turn this off to disable bitarray optimisations class _PrimeSieveE6(object): """ A prime sieve. The 'array' parameter can be used to specify a list like class to implement the sieve. Possible values for this are: list - use standard Python list bytearray - faster and uses less space (default) bitarray - (if you have it) less space that bytearray, but more time than list >>> _PrimeSieveE6(50).contents() [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47] >>> primes = _PrimeSieveE6(1000000) >>> primes.is_prime(10001) False >>> 10007 in primes True >>> sum(primes) == 37550402023 True >>> list(primes.irange(2, 47)) [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47] NOTE: if you make a large sieve it will use up lots of memory. """ # a width 6 prime sieve # # the sieve itself represents numbers with a residue of 1 and 5 modulo 6 # # i: 0 1 | 2 3 | 4 5 | 6 7 | 8 9 | 10 11 | ... # p: 1 5 | 7 11 | 13 17 | 19 23 | 25 29 | 31 35 | ... # # i->p = (3 * i) + (i & 1) + 1 # p->i = p // 3 (or in general: p->i = (p + 1) // 3 - (p % 6 == 5)) # # to check numbers up to (but not including) n we need a sieve of size: (n // 3) + (n % 6 == 2) def __init__(self, n, array=_primes_array, verbose=0): "make a sieve of primes up to " # initial sieve self.sieve = array([0]) self.max = 1 self.num = None # record the number of primes # return n copies of True or False values self.T = (lambda n, v=array([1]): v * n) # used to extend the array self.F = (lambda n, v=array([0]): v * n) # used to exclude non-primes if _primes_array_optimise: if array.__name__ == 'bitarray': # bitarray can set all values in a slice to the same boolean value if verbose: printf("[{x}: optimising for array={a}]", x=self.__class__.__name__, a=array.__name__) self.F = (lambda n: 0) # other parameters self.expandable = 0 # set to 1 to allow automatic expansion self.verbose = verbose # now extend the sieve to the required size self.extend(n) def __repr__(self): return self.__class__.__name__ + '(max=' + repr(self.max) + ')' def extend(self, n): """ extend the sieve up to (at least) """ if not (n > self.max): return #if self.verbose: printf("[{x}: expanding to {n}]", x=self.__class__.__name__) # extend the sieve to the right size s = self.sieve lo = len(s) + 1 hi = ((n + 1) // 3) + (n % 6 == 1) s.extend(self.T(hi - lo + 1)) # remove multiples of primes p from indices lo to hi for i in irange(1, isqrt(n) // 3): if s[i]: odd = (i & 1) p = (i * 3) + odd + 1 k = 2 * p # remove multiples of p starting from p^2 j = (p * p) // 3 if j < lo: j += k * ((lo - j) // k) s[j::k] = self.F((hi - j - 1) // k + 1) ##printf("eliminating {p} {ns}", ns=list((z * 3) + (z & 1) + 1 for z in irange(j, hi-1, step=k))) # remove multiples with the other residue q = p + (2 if odd else 4) j = (p * q) // 3 if j < lo: j += k * ((lo - j) // k) s[j::k] = self.F((hi - j - 1) // k + 1) ##printf("eliminating {p} {ns}", ns=list((z * 3) + (z & 1) + 1 for z in irange(j, hi-1, step=k))) self.max = n self.num = None if self.verbose: if s.__class__.__name__ == 'bitarray' and hasattr(s, 'buffer_info'): b = s.buffer_info()[1] elif hasattr(s, '__sizeof__'): b = s.__sizeof__() else: b = '' if b: b = sprintf(" ({b} bytes used)") printf("[{x}: expanded to {n}{b}]", x=self.__class__.__name__) # return the contents of the sieve (more space) [used to be called list()] def contents(self, fn=list): """ return a collection of primes in the sieve (default is a list in numerical order). this will require more memory than using generate(). """ return fn(_PrimeSieveE6.generate(self)) # return a generator (less space) def generate(self, start=0, end=None): """ generate primes in the sieve (in numerical order). the range of primes can be restricted to starting at and ending at (primes less than will be returned) (this will require less memory than contents()) """ if end is None or end > self.max: end = self.max + 1 if start < 3 and end > 2: yield 2 if start < 4 and end > 3: yield 3 s = self.sieve # generate primes from up to (but not including) for i in xrange((start + 1) // 3 - (start % 6 == 5), (end + 1) // 3 - (end % 6 == 5)): if s[i]: yield (i * 3) + (i & 1) + 1 # make this an iterable object __iter__ = generate # xrange(a, b) - generate primes in the range [a, b) - is the same as generate() now xrange = generate range = generate # irange = inclusive range def irange(self, a, b): if not (b is None or b == inf): b += 1 return self.generate(a, b) # prime test (may throw IndexError if n is too large) def is_prime(self, n): """ check to see if the number is a prime. (may throw IndexError for numbers larger than the sieve). """ if n < 2: return False # 0, 1 -> F if n < 4: return True # 2, 3 -> T r = n % 6 if r != 1 and r != 5: return False # (n % 6) != (1, 5) -> F if self.expandable: self.expand(n) return bool(self.sieve[n // 3]) prime = is_prime # allows use of "in" __contains__ = is_prime # before, after: return the prime immediately before/after n def before(self, n): """ return the largest prime less than """ if n < 3: return None if n < 4: return 2 if n < 6: return 3 if self.expandable: self.expand(n) i = (n + 1) // 3 - (n % 6 == 5) while True: i -= 1 if self.sieve[i]: return (i * 3) + (i & 1) + 1 def after(self, n): """ return the smallest prime greater than """ if n < 2: return 2 if n < 3: return 3 i = (n + 1) // 3 + (n % 6 == 1) while True: if self.expandable and not (i < len(self.sieve)): self.expand() if self.sieve[i]: return (i * 3) + (i & 1) + 1 i += 1 def between(self, a, b, fn=None): """ return primes in [a, b] """ if self.expandable: self.extend(b) r = self.irange(a, b) return (r if fn is None else fn(r)) # size = number of primes (currently) in the sieve def size(self): if self.num is None: self.num = icount(self.generate(0, self.max + 1)) return self.num __len__ = size # generate prime factors of using the sieve # (try setting mr=100 if checking large numbers) # (or mr=inf to perform all heuristic tests after the sieve is exhausted) def prime_factor(self, n, end=None, mr=0, mrr=0): """ generate (, ) pairs in the prime factorisation of positive integer , for primes in the sieve (less than ). if is set the program will use a Miller-Rabin probabilistic test after primes have failed to divide the residue to see if it is prime, and after the primes in the sieve are exhausted the Pollard Rho algorithm is used to look for remaining large prime factors. Note: By default this will only return primes up to the limit of the sieve, so may not be a complete factorisation of . However when is set it will also attempt to look for larger probabalistic prime factors. """ if n == 1: return () return prime_factor_h(n, self, end=end, nf=mr, mr=mr, mrr=mrr) # functions that can use self.prime_factor() instead of simple prime_factor() # return a list of the factors of n def factor(self, n, end=None, mr=0, mrr=0): """ return a list of the prime factors of positive integer . Note: This will only consider primes up to the limit of the sieve, this is a complete factorisation for up to the square of the limit of the sieve. """ return factor(n, fn=(lambda n: self.prime_factor(n, end=end, mr=mr, mrr=mrr))) def divisors(self, n, end=None, mr=0, mrr=0): return divisors(n, fn=(lambda n: self.prime_factor(n, end=end, mr=mr, mrr=mrr))) def divisors_pairs(self, n, end=None, mr=0, mrr=0): return divisors_pairs(n, fn=(lambda n: self.prime_factor(n, end=end, mr=mr, mrr=mrr))) def tau(self, n, end=None, mr=0, mrr=0): return tau(n, fn=(lambda n: self.prime_factor(n, end=end, mr=mr, mrr=mrr))) def is_square_free(self, n, end=None, mr=0, mrr=0): return is_square_free(n, fn=(lambda n: self.prime_factor(n, end=end, mr=mr, mrr=mrr))) # an expandable version of the sieve class _PrimeSieveE6X(_PrimeSieveE6): """ Make an expanding sieve of primes with an initial maximum of . When the sieve is expanded the function is used to calculate the new maximum, based on the previous maximum. The default function doubles the maximum at each expansion. To find the 1000th prime, (actually a list of length 1 starting with the 1000th prime): >>> primes = _PrimeSieveE6X(1000) >>> first(primes, 1, 999) [7919] We can then find the one millionth prime and the generator will expand as necessary: >>> first(primes, 1, 999999) [15485863] We can see what the current maximum number considered is: >>> primes.max 16384000 And can test for primality up to this value: >>> 1000003 in primes True The sieve will automatically expand as it is used: >>> primes.is_prime(17000023) True >>> primes.max 17000023 """ def __init__(self, n, array=_primes_array, fn=_primes_chunk, verbose=0): """ make a sieve of primes with an initial maximum of . when the sieve is expanded the function is used to calculate the new maximum, based on the previous maximum. the default function doubles the maximum at each expansion. """ _PrimeSieveE6.__init__(self, n, array=array, verbose=verbose) self.chunk = fn self.expandable = 1 # expand the sieve up to n, or by the next chunk def extend(self, n=None): """ extend the sieve to include primes up to (at least) n. if n is not specified that sieve will be expanded according to the function specified in __init__(). """ if n is None: n = self.chunk(self.max) _PrimeSieveE6.extend(self, n) return self # for backwards compatibility expand = extend # generate all primes, a chunk at a time # end = inf (or None), will expand the sieve for ever # end = self.max, will not expand the sieve def generate(self, start=0, end=inf): """ generate primes without limit, expanding the sieve as necessary. eventually the sieve will consume all available memory. """ if end is None: end = inf while start < end: # generate all primes currently in the sieve for p in _PrimeSieveE6.generate(self, start, end): yield p # expand the sieve for the next batch start = max(start, self.max + 1) if start < end: self.expand() # make this an iterable object __iter__ = generate # expand the sieve as necessary def range(self, a=0, b=None): """ generate primes in the (inclusive) range [a, b]. the sieve is expanded as necessary beforehand. """ # have we asked for unlimited generation? if b is None or b == inf: return self.generate(a) # otherwise, upper limit is provided self.extend(b) return _PrimeSieveE6.range(self, a, b) # create a suitable prime sieve def Primes(n=None, expandable=0, array=_primes_array, fn=_primes_chunk, verbose=0): """ Return a suitable prime sieve object. n - initial limit of the sieve (the sieve contains primes up to ) expandable - should the sieve expand as necessary array - list implementation to use fn - function used to increase the limit on expanding sieves If we are interested in a limited collection of primes, we can do this: >>> primes = Primes(50) >>> primes.contents() [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47] >>> sum(primes) 328 >>> 39 in primes False The collection can be extended manually to a new upper limit: >>> primes.extend(100) >>> sum(primes) 1060 >>> 97 in primes True but it doesn't automatically expand. If we want an automatically expanding version, we can set the 'expandable' flag to True. >>> primes = Primes(50, expandable=1) We can find out the current size and contents of the sieve: >>> primes.max 50 >>> primes.contents() [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47] But if we use it as a generator it will expand indefinitely, so we can only sum a restricted range: >>> sum(primes.range(0, 100)) 1060 If you don't know how many primes you'll need you can just use Primes() and get an expandable sieve with primes up to 1024, and the limit will double each time the sieve is expanded. So, to sum the first 1000 primes: >>> sum(first(primes, 1000)) 3682913 """ # if n is None then make it expandable by default if n is None: (n, expandable) = (_primes_size, True) # return an appropriate object if expandable: return _PrimeSieveE6X(n, array=array, fn=fn, verbose=verbose) else: return _PrimeSieveE6(n, array=array, verbose=verbose) # backwards compatibility def PrimesGenerator(n=None, array=_primes_array, fn=_primes_chunk): "provided for backwatds compatability. use Primes() instead." return Primes(n, expandable=1, array=array, fn=fn) # default expandable sieve primes = Primes(1, expandable=1, array=_primes_array, fn=(lambda n: _primes_size if n < _primes_size else 2 * n)) ############################################################################### # Magic Square Solver: # this is probably a bit of overkill but it works and I already had the code written class Impossible(Exception): pass class Solved(Exception): pass class MagicSquare(object): """ A magic square solver. e.g. to create a 3x3 magic square with 1 at the top centre: >>> p = MagicSquare(3) >>> p.set(1, 1) >>> p.solve() True >>> p.output() [6] [1] [8] [7] [5] [3] [2] [9] [4] When referencing the individual squares through set() and get() the squares are linearly indexed from 0 through to n^2-1. If you haven't filled out initial values for any squares then solve() is likely to take a while, especially for larger magic squares. NOTE: This class will currently find _a_ solution for the square (if it is solvable), but the solution is not necessarily unique. A future version of this code may include a function that generates all possible solutions. """ # an n x n magic square def __init__(self, n, numbers=None, lines=None): """ create an empty n x n magic square puzzle. The numbers to fill out the magic square can be specified. (If they are not specified numbers from 1 to n^2 are used). The magic lines can be specified as tuples of indices. (If they are not specified then it is assumed that all the n rows, n columns and 2 diagonals should sum to the magic value). """ n2 = n * n if numbers is None: numbers = irange(1, n2) numbers = set(numbers) s = sum(numbers) // n # make the magic lines if lines is None: lines = [] l = tuple(xrange(n)) for i in xrange(n): # row lines.append(tuple(i * n + j for j in l)) # column lines.append(tuple(i + j * n for j in l)) # diagonals lines.append(tuple(j * (n + 1) for j in l)) lines.append(tuple((j + 1) * (n - 1) for j in l)) self.n = n self.s = s self.square = [0] * n2 self.numbers = numbers self.lines = lines def set(self, i, v): """set the value of a square (linearly indexed from 0)""" self.square[i] = v self.numbers.remove(v) def get(self, i): """get the value of a square (linearly indexed from 0)""" return self.square[i] def output(self): """print the magic square""" m = max(self.square) n = (int(math.log10(m)) + 1 if m > 0 else 1) fmt = "[{:>" + str(n) + "s}]" for y in xrange(self.n): for x in xrange(self.n): v = self.square[y * self.n + x] print(fmt.format(str(v) if v > 0 else ''), end=' ') print('') # complete: complete missing squares def complete(self): """strategy to complete missing squares where there is only one possibility""" if not self.numbers: raise Solved() for line in self.lines: ns = tuple(self.square[i] for i in line) z = ns.count(0) if z == 0 and sum(ns) != self.s: raise Impossible() if z != 1: continue v = self.s - sum(ns) if v not in self.numbers: raise Impossible() i = ns.index(0) self.set(line[i], v) return True return False # hypothetical: make a guess at a square def clone(self): """return a copy of this object""" return copy.deepcopy(self) def become(self, other): """set the attributes of this object from the other object""" for x in vars(self).keys(): setattr(self, x, getattr(other, x)) def hypothetical(self): """strategy that guesses a square and sees if the puzzle can be completed""" if not self.numbers: raise Solved() i = self.square.index(0) for v in self.numbers: new = self.clone() new.set(i, v) if new.solve(): self.become(new) return True raise Impossible() # solve the square def solve(self): """solve the puzzle, returns True if a solution is found""" try: while True: while self.complete(): pass self.hypothetical() except (Impossible, Solved): pass if len(self.numbers) > 0: return False for line in self.lines: if sum(self.square[i] for i in line) != self.s: return False return True ############################################################################### # Substituted Sum Solver # originally written for Enigma 63, but applicable to lots of Enigma puzzles # a substituted sum solver # terms - list of summands of the sum (each the same length as result) # result - the result of the sum (sum of the terms) # digits - set of unallocated digits # l2d - map from letters to allocated digits # d2i - map from digits to letters that cannot be allocated to that digit # n - column we are working on (string index in result) # carry - carry from the column to the right # base - base we are working in # solutions are returned as assignments of letters to digits (the l2d dict) def _substituted_sum(terms, result, digits, l2d, d2i, n, carry=0, base=10): # are we done? if n == 0: if carry == 0: l2d.pop('_', None) yield l2d return # move on to the next column n -= 1 # find unallocated letters in this column u = list(uniq(t[n] for t in terms if t[n] not in l2d)) # and allocate them from the remaining digits for ds in itertools.permutations(digits, len(u)): _l2d = update(l2d, u, ds) # sum the column (c, r) = divmod(sum(_l2d[t[n]] for t in terms) + carry, base) # is the result what we expect? if result[n] in _l2d: # the digit of the result is already allocated, check it if _l2d[result[n]] != r: continue allocated = ds else: # the digit in the result is one we haven't come across before if r not in digits or r in ds: continue _l2d[result[n]] = r allocated = ds + (r,) # check there are no invalid allocations if any(any(_l2d[x] == d for x in ls if x in _l2d) for (d, ls) in d2i.items()): continue # try the next column for r in _substituted_sum(terms, result, digits.difference(allocated), _l2d, d2i, n, c, base): yield r def substitute(s2d, text, digits=None): """ given a symbol-to-digit mapping and some text , return the text with the digits (as defined by the sequence ) substituted for the symbols. characters in the text that don't occur in the mapping are unaltered. if there are braces present in then only those portions of the enclosed in braces are substituted. >>> substitute(dict(zip('DEMNORSY', (7, 5, 1, 6, 0, 8, 9, 2))), "SEND + MORE = MONEY") '9567 + 1085 = 10652' """ if text is None: return None if digits is None: digits = base_digits() return translate(text, (lambda x: digits[s2d[x]] if x in s2d else x)) # friendly interface to the substituted sum solver def substituted_sum(terms, result, digits=None, l2d=None, d2i=None, base=10): """ a substituted addition sum solver - encapsulated by the SubstitutedSum class. terms - list of summands in the sum result - result of the sum (sum of the terms) digits - digits to be allocated (default: 0 - base-1, less any allocated digits) l2d - initial allocation of digits (default: all digits unallocated) d2i - invalid allocations (default: leading digits cannot be 0) base - base we're working in (default: 10) """ # fill out the parameters if l2d is None: l2d = dict() if digits is None: digits = xrange(base) digits = set(digits).difference(l2d.values()) if d2i is None: d2i = dict() d2i[0] = set(x[0] for x in itertools.chain(terms, [result]) if len(x) > 1) # number of columns in sum n = len(result) # make sure the terms are the same length as the result ts = list('_' * (n - len(t)) + t for t in terms) assert all(len(t) == n for t in ts), "result is shorter than terms" l2d['_'] = 0 # call the solver for r in _substituted_sum(ts, result, digits, l2d, d2i, n, 0, base): yield r # object interface to the substituted sum solver class SubstitutedSum(object): """ A solver for addition sums with letters substituted for digits. A substituted sum object has the following useful attributes: terms - the summands of the sum result - the result of the sum text - a textual form of the sum (e.g. " + = ") e.g. Enigma 21: "Solve: BPPDQPC + PRPDQBD = XVDWTRC" >>> SubstitutedSum(['BPPDQPC', 'PRPDQBD'], 'XVDWTRC').go() BPPDQPC + PRPDQBD = XVDWTRC 3550657 + 5850630 = 9401287 / B=3 C=7 D=0 P=5 Q=6 R=8 T=2 V=4 W=1 X=9 e.g. Enigma 63: "Solve: LBRLQQR + LBBBESL + LBRERQR + LBBBEVR = BBEKVMGL" >>> SubstitutedSum(['LBRLQQR', 'LBBBESL', 'LBRERQR', 'LBBBEVR'], 'BBEKVMGL').go() LBRLQQR + LBBBESL + LBRERQR + LBBBEVR = BBEKVMGL 8308440 + 8333218 + 8302040 + 8333260 = 33276958 / B=3 E=2 G=5 K=7 L=8 M=9 Q=4 R=0 S=1 V=6 """ def __init__(self, terms, result, base=10, digits=None, l2d=None, d2i=None): """ create a substituted addition sum puzzle. terms - a list of the summands of the sum. result - the result of the sum (i.e. the sum of the terms). The following parameters are optional: base - the number base the sum is in (default: 10) digits - permissible digits to be substituted in (default: determined from base) l2d - initial map of letters to digits (default: all letters unassigned) d2i - map of digits to invalid letter assignments (default: leading digits cannot be 0) If you want to allow leading digits to be 0 pass an empty dictionary for d2i. """ text = join(terms, sep=' + ') + ' = ' + result self.terms = terms self.base = base self.result = result self.digits = digits self.l2d = l2d self.d2i = d2i self.text = text def solve(self, check=None, verbose=0): """ generate solutions to the substituted addition sum puzzle. solutions are returned as a dictionary assigning letters to digits. """ if verbose > 0: printf("{self.text}") for s in substituted_sum(self.terms, self.result, base=self.base, digits=self.digits, l2d=self.l2d, d2i=self.d2i): if check and (not check(s)): continue if verbose > 0: self.output_solution(s) yield s def substitute(self, s, text, digits=None): """ given a solution to the substituted sum and some text return the text with letters substituted for digits. """ return substitute(s, text, digits=digits) def output_solution(self, s, digits=None): """ given a solution to the substituted sum output the assignment of letters to digits and the sum with digits substituted for letters. """ printf("{t} / {s}", # print the sum with digits substituted in t=substitute(s, self.text, digits=digits), # output the assignments in letter order s=map2str(s, sep=" ", enc="") ) solution = output_solution def run(self, check=None, first=0): """find all solutions (matching the filter ) and output them""" for s in self.solve(check=check, verbose=1): if first: break # backwards compatability go = run # class method to chain multiple sums together @classmethod def chain(cls, sums, base=10, digits=None, l2d=None, d2i=None): """ solve a sequence of substituted sum problems. sums are specified as a sequence of: (, , ..., ) """ # are we done? if not sums: yield l2d else: # solve the first sum s = sums[0] for r in cls(s[:-1], s[-1], base=base, digits=digits, l2d=l2d, d2i=d2i).solve(): # and recursively solve the rest for x in cls.chain(sums[1:], base=base, digits=digits, l2d=r, d2i=d2i): yield x @classmethod def chain_run(cls, sums, base=10, digits=None, l2d=None, d2i=None): template = join(('(' + join(s[:-1], sep=' + ') + ' = ' + s[-1] + ')' for s in sums), sep=' ') printf("{template}") for s in cls.chain(sums, base=base, digits=digits, l2d=l2d, d2i=d2i): printf("{t} / {s}", t=substitute(s, template), s=map2str(s, sep=" ", enc="") ) # backwards compatability chain_go = chain_run # class method to call from the command line @classmethod def run_command_line(cls, args): """ run the SubstitutedSum solver with the specified command line arguments. e.g. Enigma 327 % python enigma.py SubstitutedSum "KBKGEQD + GAGEEYQ + ADKGEDY = EXYAAEE" (KBKGEQD + GAGEEYQ + ADKGEDY = EXYAAEE) (1912803 + 2428850 + 4312835 = 8654488) / A=4 B=9 D=3 E=8 G=2 K=1 Q=0 X=6 Y=5 Optional parameters: --base= (or -b) Specifies the numerical base of the sum, solutions are printed in the specified base (so that the letters and digits match up, but don't get confused by digits that are represented by letters in bases >10). e.g. Enigma 1663 % python enigma.py SubstitutedSum --base=11 "FAREWELL + FREDALO = FLINTOFF" (FAREWELL + FREDALO = FLINTOFF) (61573788 + 657A189 = 68042966) / A=1 D=10 E=7 F=6 I=0 L=8 N=4 O=9 R=5 T=2 W=3 (6157A788 + 6573189 = 68042966) / A=1 D=3 E=7 F=6 I=0 L=8 N=4 O=9 R=5 T=2 W=10 --assign=, (or -a,) Assign a digit to a letter. There can be multiple --assign options. e.g. Enigma 1361 % python enigma.py SubstitutedSum --assign="O,0" "ELGAR + ENIGMA = NIMROD" (ELGAR + ENIGMA = NIMROD) (71439 + 785463 = 856902) / A=3 D=2 E=7 G=4 I=5 L=1 M=6 N=8 O=0 R=9 --digits=,,... (or -d,,...) Specify the digits that can be assigned to unassigned letters. (Note that the values of the digits are specified in base 10, even if you have specified a --base= option) e.g. Enigma 1272 % python enigma.py SubstitutedSum --digits="0-8" "WILKI + NSON = JONNY" (WILKI + NSON = JONNY) (48608 + 3723 = 52331) / I=8 J=5 K=0 L=6 N=3 O=2 S=7 W=4 Y=1 (48708 + 3623 = 52331) / I=8 J=5 K=0 L=7 N=3 O=2 S=6 W=4 Y=1 --invalid=, (or -i,) Specify letters that cannot be assigned to a digit. There can be multiple --invalid options, but they should be for different s. If there are no --invalid options the default will be that leading zeros are not allowed. If you want to allow them you can give a "--invalid=0," option. Enigma 171 % python enigma.py SubstitutedSum -i0,016 -i1,1 -i3,3 -i5,5 -i6,6 -i7,7 -i8,8 -i9,9 "1939 + 1079 = 6856" (1939 + 1079 = 6856) (2767 + 2137 = 4904) / 0=1 1=2 3=6 5=0 6=4 7=3 8=9 9=7 --help (or -h) Prints a usage summary. """ usage = join(( sprintf("usage: {cls.__name__} [] \" + + ... = \" ..."), "options:", " --base= (or -b) = set base to ", " --assign=, (or -a,) = assign digit to letter", " --digits=,... or - (or -d...) = available digits", " --invalid=, (or -i,) = invalid digit to letter assignments", " --help (or -h) = show command-line usage", ), sep=nl) # process options (--[=] or -) opt = dict(l2d=dict(), d2i=None) while args and args[0].startswith('-'): arg = args.pop(0) try: if arg.startswith('--'): (k, _, v) = arg.lstrip('-').partition('=') else: (k, v) = (arg[1], arg[2:]) if k == 'h' or k == 'help': print(usage) return -1 elif k == 'b' or k == 'base': # --base= (or -b) opt['base'] = int(v) elif k == 'a' or k == 'assign': # --assign=, (or -a) (l, d) = v.split(',', 1) opt['l2d'][l] = int(d) elif k == 'd' or k == 'digits': # --digits=,... or - (or -d) opt['digits'] = _digits(v) elif k == 'i' or k == 'invalid': # --invalid=, (or -i,) if opt['d2i'] is None: opt['d2i'] = dict() (ds, s) = _split(v, maxsplit=-1) for i in _digits(ds): opt['d2i'][i] = opt['d2i'].get(i, set()).union(s) else: raise ValueError() except Exception: printf("{cls.__name__}: invalid option: {arg}") return -1 # check command line usage if not args: print(usage) return -1 # extract the sums sums = list(re.split(r'[\s\+\=]+', arg) for arg in args) # call the solver cls.chain_go(sums, **opt) return 0 ############################################################################### # Generic Substituted Expression Solver # be aware that the generated code can fall foul of restrictions in the Python # interpreter. # # the standard Python interpreter will throw: # # "SystemError: too many statically nested blocks"; or: # "SyntaxError: too many statically nested blocks" # # if there are more than 20 levels of nested loops. # # and it will also throw: # # "IndentationError: too many levels of indentation" # # if there are more than 100 levels of indentation. # # the PyPy interpreter has neither of these limitations # # the experimental [[ denest=1 ]] parameter will produce less nested code # to work around this issue. # TODO: think about negative values - expressions resulting in an # alphametic word must be non-negative # # TODO: consider ordering the symbols, so we can calculate words sooner. # # TODO: consider allowing a "wildcard" character, for symbols that can take # on any available digit (but still not allow leading zeros). [Enigma 1579] # # TODO: consider allowing code to generate possible values for a word # like: "[possible values] -> WORD", (so [Teaser 3019] would be: # [ "primes.irange(1000, 9999) -> WILL", "primes.irange(10, 99) -> AM" ], or # [Teaser 3018] could use [ "powers(31, 99, 2) -> ABCD" ]) # # TODO: spotting expressions for independent groups and solving # each group separately [Teaser 2990] _SYMBOLS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' # find words in string def _find_words(s, r=1): words = set(re.findall(r'{(\w+?)}', s)) if r: # return the words return words else: # return individual characters return set().union(*words) # replace words in string def _replace_words(s, fn): # new style, with braces _fn = lambda m: fn(m.group(1)) return re.sub(r'{(\w+?)}', _fn, s) # local variable used to represent symbol x _sym = lambda x: '_' + x # return an expression that evaluates word in base def _word(w, base): (m, d) = (1, dict()) for x in w[::-1]: d[x] = d.get(x, 0) + m m *= base return join((concat((_sym(k),) + (() if v == 1 else ('*', v))) for (k, v) in d.items()), sep=' + ') @static(i=0) def gensym(x): """ generate a unique string starting with . >> gensym('foo') 'foo1' >> gensym('foo') 'foo2' """ gensym.i += 1 return concat(x, gensym.i) # file.writelines does NOT include newline characters def writelines(fh, lines, sep=None, flush=1): if sep is None: sep = os.linesep for line in lines: fh.write(line) fh.write(sep) if flush: fh.flush() # split string on any of the characters in _split_sep = ',|+' def _split(s, sep=_split_sep, maxsplit=0): d = sep[0] if len(sep) > 1: # map all other separators to d s = re.sub(r'[' + re.escape(sep[1:]) + r']', d, s) if maxsplit == 0: return s.split(d) elif maxsplit > 0: return s.split(d, maxsplit) else: return s.rsplit(d, -maxsplit) # a sequence of digit values may be specified (in decimal) as: # "-" = a range of digits # ",...," "|...|" "+...+" # "" = a single digit # returns a sequence of integers def _digits(s): # "-" if '-' in s: (a, _, b) = s.partition('-') return irange(int(a), int(b)) # ",...," "|...|" "+...+" # "" return tuple(int(d) for d in _split(s, _split_sep)) # fix up implicit parameters # if contains no {word} parameters, then enclose words from def _fix_implicit(s, symbols): if s is None: return None if re.search(r'{\w+}', s): return s # was: [[ if '{' in s: return s ]] return re.sub(r'[' + symbols + r']+', (lambda m: '{' + m.group(0) + '}'), s) def _fix_implicit_seq(seq, symbols): if seq is None: return None return list(_fix_implicit(x, symbols) for x in seq) class SubstitutedExpression(object): """ A solver for Python expressions with symbols substituted for numbers. It takes a Python expression and then tries all possible ways of assigning symbols (by default the capital letters) in it to digits and returns those assignments which result in the expression having a True value. This allows for more general expressions to be evaluated than specialised solvers, like SubstitutedSum(), alow. Enigma 1530 >>> SubstitutedExpression('TOM * 13 = DALEY').run().n (TOM * 13 = DALEY) (796 * 13 = 10348) / A=0 D=1 E=4 L=3 M=6 O=9 T=7 Y=8 [1 solution] 1 See SubstitutedExpression.run_command_line() for more examples. """ @classmethod def set_default(cls, **kw): """ set default values for instance initialisation. """ for (k, v) in kw.items(): cls.defaults[k] = v # standard default values # add new parameters here, and they will be set up automatically # but you still need to document them in __init__ defaults = dict( # parameters that have a default value base=10, distinct=1, process=1, reorder=1, first=0, denest=0, sane=1, verbose=1, # other parameters exprs=None, symbols=None, digits=None, s2d=None, d2i=None, answer=None, accumulate=None, literal=None, template=None, solution=None, header=None, check=None, env=None, code=None, ) def __init__(self, exprs, base=None, symbols=None, digits=None, s2d=None, l2d=None, d2i=None, answer=None, accumulate=None, literal=None, template=None, solution=None, header=None, distinct=None, check=None, env=None, code=None, process=None, reorder=None, first=None, denest=None, sane=None, verbose=None ): """ create a substituted expression solver. exprs - the expression(s) exprs can be a single expression, or a sequence of expressions. A single expression is of the form: "" or " = " or (, ) where value is a valid "word" (sequence of symbols), or an integer value. The following parameters are optional: base - the number base to operate in (default: 10) symbols - the symbols to substituted in the expression (default: upper case letters) digits - the digits to be substituted in (default: determined from base) s2d - initial map of symbols to digits (default: all symbols unassigned) d2i - map of digits to invalid symbol assignments (default: leading digits cannot be 0) distinct - symbols which should have distinct values (1 = all, 0 = none) (default: 1) literal - symbols which stand for themselves (e.g. "012") (default: None) answer - an expression for the answer value accumulate - accumulate answers using the specified object check - a boolean function used to accept/reject solutions (default: None) env - additional environment for evaluation (default: None) code - additional lines of code evaluated before solving (default: None) denest - work around CPython statically nested block limit sane - enable/disable sanity checks (default: 1) verbose - control informational output (default: 1) If you want to allow leading digits to be 0 pass an empty dictionary for d2i. """ # return the first not-None value, or defaults[key] def get_default(key, *values): for v in values: if v is not None: return v return self.__class__.defaults.get(key, None) # set defaults from class, most are simple except: # s2d - use l2d if s2d is not specified (for backward compatability) scope = locals() for k in self.__class__.defaults.keys(): if k == 's2d': v = get_default(k, s2d, l2d) else: v = get_default(k, scope[k]) setattr(self, k, v) self._processed = 0 # set by process self._prepared = 0 # set by prepare() if self.process: self._process() # verbose flags: # generate output vH = 4 # output header template vT = 8 # output solutions (from template) #vS = 512 # output solution symbol -> digit mapping vA = 16 # output answer counts / accumulated measures # information / debugging vE = 32 # output elapsed time vP = 64 # output solver parameters vI = 128 # output solver info vC = 256 # output code (before compilation) # standard debug levels v1 = vH | vT | vA # 1 = header + solutions + count v2 = v1 | vI # 2 = 1 + solver info v3 = v2 | vE | vC # 3 = 2 + timing + code v9 = vH | vT | vA | vE | vP | vI | vC # 9 = everything def _verbose(self, n): if not n: return 0 # if it is a string, parse it into a number if isinstance(n, basestring): # sort out "old-style" numeric arguments if re.match(r'[\d\-\|\+\,]+', n): n = sum(_digits(n)) else: # sort out "new style" flags # "1+E" "1-T" "1-T+E" "9" "HTAEPIC" assert False # old style verbose flags (1, 2, 3) if n < 4: return (0, self.v1, self.v2, self.v3)[n] # otherwise return n # sort out calling methods def _process(self): exprs = self.exprs base = self.base symbols = self.symbols digits = self.digits s2d = self.s2d d2i = self.d2i answer = self.answer template = self.template distinct = self.distinct literal = self.literal denest = self.denest process = self.process sane = self.sane verbose = self.verbose # sort out verbose argument verbose = self._verbose(verbose) if sane == 0 and verbose > 0: printf("WARNING: sanity checks disabled - good luck!") # the symbols to replace (for implicit expressions) if symbols is None: symbols = _SYMBOLS # process expr to be a list of (, ) pairs, where: # is: # None = look for a true value # word = look for a value equal to the substituted word # integer = look for the specific value if process: # allow expr to be a single string if isinstance(exprs, basestring): exprs = [exprs] # now process the list xs = list() for expr in exprs: if isinstance(expr, basestring): # expression is a single string, turn it into an (, ) pair (v, s) = ('', re.split(r'\s+=\s+', expr)) if len(s) == 2: (expr, v) = s if not v: v = None else: # assume expr is already an (, ) pair (expr, v) = expr # convert implicit (without braces) into explicit (with braces) if symbols: if isinstance(v, basestring) and '{' not in v and '{' not in expr and all(x in symbols for x in v): v = '{' + v + '}' expr = _fix_implicit(expr, symbols) # value is either an alphabetic or a numeric literal if isinstance(v, basestring) and '{' not in v: v = base2int(v, base=base) xs.append((expr, v)) exprs = xs # fix up implicit (old style) parameters answer = _fix_implicit(answer, symbols) template = _fix_implicit(template, symbols) # make the output template (which is kept in input order) # and also categorise the expressions to (, , ), where is: # 0 = answer ( = None) # 1 = expression with no value, we look for a true ( = None) # 2 = expression with an integer value, we do a direct comparison ( = int) # 3 = expression with string value, we look to assign/check symbols ( = str) ts = list() xs = list() for (x, v) in exprs: if v is None: ts.append(x) xs.append((x, v, 1)) elif isinstance(v, basestring): ts.append(x + ' = ' + v) xs.append((x, v[1:-1], 3)) else: ts.append(x + ' = ' + int2base(v, base=base)) xs.append((x, v, 2)) if answer: ts.append(answer) xs.append((answer, None, 0)) _template = join(('(' + t + ')' for t in ts), sep=' ') exprs = xs # initial mapping of symbols to digits if s2d is None: s2d = dict() # literal values are symbols which stand for themselves if literal is None: literal = () for s in literal: s2d[s] = base2int(s, base=base) # allowable digits (and invalid digits) if digits is None: digits = set(xrange(base)) else: digits = set(digits) if sane > 0: ds = set(xrange(base)) if verbose > 0: # check for invalid digits for d in digits: if d not in ds: printf("WARNING: SubstitutedExpression: non-valid digit {d} for base {base} specified", d=repr(d)) digits.intersection_update(ds) # TODO: I suspect this needs to work with more values of "distinct" # if all values are distinct (including literals), remove them from digits if distinct == 1: digits = digits.difference(s2d.values()) idigits = set(xrange(base)).difference(digits) # find words in all exprs words = _find_words(_template) # and determine the symbols that are used symbols = join(sorted(set().union(*words))) # invalid (, ) assignments invalid = set() if d2i is not None: # it should provide a sequence of (, ) pairs for (d, ss) in (d2i.items() if hasattr(d2i, 'items') else d2i): if sane > 0 and verbose > 0 and d not in digits and d not in idigits: printf("WARNING: SubstitutedExpression: non-valid invalid digit {d} specified", d=repr(d)) invalid.update((s, d) for s in ss) else: # disallow leading zeros if 0 in digits: for w in words: if len(w) > 1: invalid.add((w[0], 0)) # but for the rest of the time we are only interested # in words in the (not the ) words = set() for (x, _, _) in exprs: words.update(w for w in _find_words(x) if len(w) > 1) # find the symbols in the (, ) pairs # xs = symbols in # vs = symbols in (xs, vs) = (list(), list()) for (x, v, k) in exprs: xs.append(_find_words(x, r=0)) vs.append(set(v) if k == 3 else set()) # determine the symbols in each expression syms = list(x.union(v) for (x, v) in zip(xs, vs)) # sort out distinct=0,1 if isinstance(distinct, int): distinct = (symbols if distinct else '') # distinct should be a sequence (probably of strings) if isinstance(distinct, basestring): distinct = [distinct] # add the value of the symbols into the template self.template = (_template if template is None else template) if self.solution is None: self.solution = join(diff(symbols, literal)) if self.header is None: self.header = _replace_words(self.template, identity) # sort out negative value for denest, will not be enabled if running under PyPy if denest < 0: denest = (0 if _pypy else -denest) # sort out denest=1 if denest == 1: denest = 50 # update the processed values self.exprs = exprs self.symbols = symbols self.digits = digits self.s2d = s2d self.answer = answer self.distinct = distinct self.literal = literal self.denest = denest self.verbose = verbose self._words = words self._invalid = invalid self._idigits = idigits self._exprs = (exprs, xs, vs, ts, syms) self._processed = 1 # create and compile the code # NOTE: the generated code can have more than 20 nested blocks, # which raises a SyntaxError in CPython. A workaround is to use PyPy # instead (which doesn't seem to have this limitation) def _prepare(self): base = self.base symbols = self.symbols digits = self.digits s2d = self.s2d d2i = self.d2i answer = self.answer distinct = self.distinct literal = self.literal env = self.env code = self.code reorder = self.reorder denest = self.denest sane = self.sane verbose = self.verbose words = self._words invalid = self._invalid idigits = self._idigits (exprs, xs, vs, ts, syms) = self._exprs # output run parameters if self.verbose & self.vP: print("-- [code] --" + nl + join(self.save(quote=1), sep=nl) + nl + "-- [/code] --") # # remove assigned symbols from distinct groups [suggested by Frits] # if s2d: # # update invalid pairs # for (k, v) in s2d.items(): # for ds in distinct: # if k in ds: # invalid.update((d, v) for d in ds if d != k) # # update distinct # distinct = list(join(x) for x in (diff(x, s2d.keys()) for x in distinct) if len(x) > 1) # valid digits for each symbol valid = dict() for s in symbols: if s in s2d: continue valid[s] = list(digits.difference(d for (x, d) in invalid if x == s)) #for k in sorted(valid.keys()): printf("{k} -> {v}", v=valid[k]) # at this point we can apply some heuristic re-writing rules: # word = value -> value = word, if value is free of alphametic symbols if sane > 0: for (i, (expr, val, k)) in enumerate(exprs): #printf("[{i}] ({expr!r}, {val!r}, {k}) xs={xs} vs={vs}", xs=xs[i], vs=vs[i]) # " == " and expr contains no alphametic symbols: # (" == ", None, 1) --> ("", "", 3) if k == 1: word = xpr = None m = re.match(r'\s*\{([' + symbols + r']+)\}\s*==\s*(.+)\s*$', expr) if m: (word, xpr) = m.groups() else: # try: " == " m = re.match(r'\s*(.+?)\s*==\s*\{([' + symbols + r']+)\}\s*$', expr) if m: (xpr, word) = m.groups() if word and expr: if re.search(r'\{[' + symbols + r']+\}', xpr) is None: if verbose > 0: printf("[SubstitutedExpression: replacing ({t}) -> ({xpr} = {word})]", t=ts[i]) exprs[i] = (xpr, word, 3) (xs[i], vs[i]) = (vs[i], xs[i]) # " = ": (even more efficient to use --assign instead) # ("", , 2) --> ("", "", 3) if k == 2: if expr[0] == '{' and expr[-1] == '}': word = expr[1:-1] if all(x in symbols for x in word): if verbose > 0: printf("[SubstitutedExpression: replacing ({t}) -> ({val} = {{{word}}})]", t=ts[i]) exprs[i] = (int2base(val, base=10), word, 3) (xs[i], vs[i]) = (vs[i], xs[i]) # reorder the expressions into a more appropriate evaluation order if reorder: # at each stage chose the expression with the fewest remaining possibilities d = set(s2d.keys()) (s, r) = (list(), list(i for (i, _) in enumerate(syms))) # formerly we used: # # ( is answer? ) (# of unassiged symbols) -(number of new symbols we get) #fn = lambda i: (exprs[i][2] == 0, len(xs[i].difference(d)), -len(vs[i].difference(d, xs[i]))) # # now we use: # # ( is answer? ) ( total possibilities for unassigned symbols ) -(number of new symbols we get) fn = lambda i: (exprs[i][2] == 0, multiply(len(valid[x]) for x in xs[i] if x not in d), -len(vs[i].difference(d, xs[i]))) while r: i = min(r, key=fn) s.append(i) d.update(xs[i], vs[i]) r.remove(i) # update the lists exprs = list(exprs[i] for i in s) xs = list(xs[i] for i in s) vs = list(vs[i] for i in s) ts = list(ts[i] for i in s) if verbose & self.vI: # output solver information printf("[base={base}, digits={digits}, symbols={symbols!r}, distinct={distinct!r}]", distinct=join(distinct, sep=',')) printf("[s2d={s2d}, d2i={d2i}]") # output the solving strategy (ss, d) = (list(), set(s2d.keys())) for (i, x) in enumerate(xs): ss.append(sprintf("({e}) [{n}+{m}]", e=ts[i], n=len(x.difference(d)), m=len(vs[i].difference(d, x)))) d.update(x, vs[i]) printf("[strategy: {ss}]", ss=join(ss, sep=' -> ')) # turn distinct into a dict mapping -> if not isinstance(distinct, dict): d = dict() for ss in distinct: if sane > 0 and verbose > 0 and len(set(ss).difference(s2d.keys())) > len(digits): printf("[SubstitutedExpression: WARNING: distinct=\"{ss}\" has more symbols than available digits]") for s in ss: if s not in d: d[s] = set() d[s].update(x for x in ss if x != s) distinct = d # generate the program (line by line) (prog, _, indent) = ([], "", " ") (vx, vy, vr) = ("_x_", "_y_", "_r_") # local variables (that don't clash with _sym(x)) # start with any initialisation code if code: # code should be a sequence (of strings) if isinstance(code, basestring): code = [code] prog.extend(code) # wrap it all up as function solver solver = gensym('_substituted_expression_solver') prog.append(sprintf("{_}def {solver}():")) _ += indent # set initial values done = set() for (s, d) in s2d.items(): prog.append(sprintf("{_}{s} = {d}", s=_sym(s))) done.add(s) # [denest] workaround statically nested block limit if denest: # set other initial values and words to None for s in symbols: if s not in s2d: prog.append(sprintf("{_}{s} = None", s=_sym(s))) for w in words: prog.append(sprintf("{_}{w} = None", w=_sym(w))) # keep track of nested functions blocks = [ gensym('_substituted_expression_block') ] block = None block_args = join(map(_sym, chain(symbols, words)), sep=", ") indent_reset = indent # look for words which can be made for w in words: if all(x in done for x in w): prog.append(sprintf("{_}{w} = {x}", w=_sym(w), x=_word(w, base))) in_loop = False # use_sets = 1 # using sets is slower # deal with each , pair for ((expr, val, k), xsyms, vsyms) in zip(exprs, xs, vs): # [denest] work around statically nested block limit if denest and block is None: # start a new function block block = blocks[-1] _ = indent_reset # In Python3 we can use [[ nonlocal ]] instead of passing the symbols around prog.append(sprintf("{_}def {block}({block_args}):")) _ += indent in_loop = False # EXPERIMENTAL: do something about: ": = if k == 3 and expr.endswith(':'): prog.append(sprintf("{_}for {vx} in {expr}")) # expr already has a colon _ += indent #prog.append(sprintf("{_}{w} = {vx}", w=_sym(val))) done.update(xsyms) else: # deal with each symbol in # TODO: we could consider these in an order that makes words # in as soon as possible for s in xsyms: if s in done: continue # allowable digits for s ds = valid[s] in_loop = True prog.append(sprintf("{_}for {s} in {ds}:", s=_sym(s))) _ += indent if done and s in distinct: # TODO: we should exclude initial values (that are excluded from ds) here check = list(_sym(x) for x in done if x in distinct[s]) if check: #if use_sets: # if len(check) == 1: # check = _sym(s) + " != " + check[0] # else: # check = _sym(s) + " not in " + join(check, sep=", ", enc="{}") #else: check = join(((_sym(s) + ' != ' + x) for x in check), sep=' and ') prog.append(sprintf("{_}if {check}:")) _ += indent done.add(s) # look for words which can now be made for w in words: if s in w and all(x in done for x in w): prog.append(sprintf("{_}{w} = {x}", w=_sym(w), x=_word(w, base))) # calculate the expression if k != 0 and (not expr.endswith(':')): # (but not for the answer expression) x = _replace_words(expr, (lambda w: '(' + _sym(w) + ')')) prog.append(sprintf("{_}try:")) prog.append(sprintf("{_} {vx} = int({x})")) prog.append(sprintf("{_}except NameError:")) # catch undefined functions prog.append(sprintf("{_} raise")) prog.append(sprintf("{_}except Exception:")) # maybe "except (ArithmeticError, ValueError)" prog.append(sprintf("{_} {skip}", skip=('continue' if in_loop else 'return'))) # check the value if k == 3: # this is a literal (alphametic) word # so it must have a non-negative value prog.append(sprintf("{_}if {vx} >= 0:")) _ += indent for (j, y) in enumerate(val[::-1], start=-len(val)): if y in done: # this is a symbol with an assigned value prog.append(sprintf("{_}{vy} = {vx} % {base}")) # check the value prog.append(sprintf("{_}if {vy} == {y}:", y=_sym(y))) _ += indent prog.append(sprintf("{_}{vx} //= {base}")) # and check x == 0 for the final value if j == -1: prog.append(sprintf("{_}if {vx} == 0:")) _ += indent else: # this is a new symbol... prog.append(sprintf("{_}{y} = {vx} % {base}", y=_sym(y))) check = list() # check it is different from existing symbols if y in distinct: check.extend(_sym(x) for x in done if x in distinct[y]) # check any invalid values for this symbol for v in idigits.union(v for (s, v) in invalid if y == s): check.append(str(v)) if check: #if use_sets: # if len(check) == 1: # check = _sym(y) + " != " + check[0] # else: # check = _sym(y) + " not in " + join(check, sep=", ", enc="{}") #else: check = join((_sym(y) + " != " + x for x in check), sep=" and ") prog.append(sprintf("{_}if {check}:")) _ += indent prog.append(sprintf("{_}{vx} //= {base}")) # and check x == 0 for the final value if j == -1: prog.append(sprintf("{_}if {vx} == 0:")) _ += indent done.add(y) # look for words which can now be made for w in words: if y in w and all(x in done for x in w): prog.append(sprintf("{_}{w} = {x}", w=_sym(w), x=_word(w, base))) elif k == 1: # look for a True value prog.append(sprintf("{_}if {vx}:")) _ += indent elif k == 2: # it's a comparable value prog.append(sprintf("{_}if {vx} == {val}:")) _ += indent # [denest] work around statically nested block limit if denest and len(_) > denest: # chain into the next block block = gensym('_substituted_expression_block') blocks.append(block) # return the current state of the symbols # Python3 can use [[ yield from ... ]] prog.append(sprintf("{_}for {vr} in {block}({block_args}): yield {vr}")) block = None # [denest] work around statically nested block limit if denest: if block is None: # we need a final trivial block block = blocks[-1] _ = indent_reset prog.append(sprintf("{_}def {block}({block_args}):")) _ += indent # close final function block prog.append(sprintf("{_}yield [{block_args}]")) _ = indent_reset # now call the first block block = blocks[0] prog.append(sprintf("{_}for [{block_args}] in {block}({block_args}):")) _ += indent # yield solutions as dictionaries d = join((("'" + s + "': " + _sym(s)) for s in sorted(done)), sep=', ') if answer: # compute the answer r = _replace_words(answer, (lambda w: '(' + _sym(w) + ')')) prog.append(sprintf("{_}{vr} = {r}")) prog.append(sprintf("{_}yield ({{ {d} }}, {vr})")) else: prog.append(sprintf("{_}yield {{ {d} }}")) # turn the program lines into a string prog = join(prog, sep=nl) if verbose & self.vC: printf("-- [code language=\"python\"] --{nl}{prog}{nl}-- [/code] --") # compile the solver # a bit of jiggery pokery to make this work in several Python versions # older Python barfs on: # ns = dict() # eval(prog, None, ns) # solve = ns[solver] if not env: env = dict() gs = update(globals(), env) try: code = compile(prog, '', 'exec') except Exception: # the program failed to compile # this can be because the supplied expressions do not form valid Python # or due to an issue in the Python interpreter itself # (e.g. in standard Python you can't have more than 20 nested blocks, # or more than 100 indent levels - PyPy does not have these limitations) printf("SubstitutedExpression: compilation error from Python interpreter [{sys.executable}]") if not (verbose & self.vC): printf("(use verbose level 256 to output code before compilation)") printf("(or use the \"denest=1\" option (--denest, -X) to reduce program complexity)") raise eval(code, gs) self._solver = gs[solver] self._globals = gs self._prepared = 1 # execute the code def solve(self, check=None, first=None, verbose=None): """ generate solutions to the substituted expression problem. solutions are returned as a dictionary assigning symbols to digits. check - a boolean function called to reject unwanted solutions first - if set to positive only the first solutions are returned verbose - if set to >0 solutions are output as they are found, >1 additional information is output. """ if not self._prepared: self._prepare() solver = self._solver answer = self.answer header = self.header if check is None: check = self.check if first is None: first = self.first verbose = (self.verbose if verbose is None else self._verbose(verbose)) if verbose & self.vH and header: printf("{header}") n = 0 for s in solver(): if check and (not check(s)): continue if verbose & self.vT: self.output_solution((s[0] if answer else s)) # return the result yield s n += 1 if first and first == n: break # output a solution as: "

use: icount(i, p, n + 1) < n + 1 This is what icount_at_most(i, p, n) does. The following will stop testing at 73 (the 21st prime): >>> icount_at_most(irange(1, 100), is_prime, 20) False If p is not specified a function that always returns True is used, so you can use this function to count the number of items in a (finite) iterator: >>> icount(Primes(1000)) 168 """ if p is None: if t is None: if hasattr(i, '__len__'): return len(i) else: # a quick way to count an iterable d = collections.deque(enumerate(i, start=1), maxlen=1) return (d[0][0] if d else 0) else: p = true n = 0 for x in i: if p(x): n += 1 if n == t: break return n # icount recipes icount_exactly = lambda i, p=None, n=None: icount(i, p, n + 1) == n icount_at_least = lambda i, p=None, n=None: icount(i, p, n) == n icount_at_most = lambda i, p=None, n=None: icount(i, p, n + 1) < n + 1 # find: like index(), but return -1 instead of throwing an error def find(seq, v): """ find the first index of a value in a sequence, return -1 if not found. for a string it works like str.find() (for single characters) >>> find('abc', 'b') 1 >>> find('abc', 'z') -1 but it also works on lists >>> find([1, 2, 3], 2) 1 >>> find([1, 2, 3], 7) -1 and on iterators in general (don't try this with a non-prime value) >>> find(primes, 10007) 1229 Note that this function works by attempting to use the index() method of the sequence. If it implements index() in a non-compatible way this function won't work. """ try: return seq.index(v) except ValueError: return -1 except AttributeError: pass if isinstance(seq, dict): # search the keys # (or we could use find() in the values, and return the correspond index in keys) for (k, x) in s.items(): if x == v: return k else: # search the sequence for (i, x) in enumerate(seq): if x == v: return i # not found return -1 def rfind(seq, v): """find the last index of a value in a sequence, return -1 if not found""" i = find(seq[::-1], v) return (-1 if i == -1 else len(s) - i - 1) # trim elements from a sequence def trim(seq, head=0, tail=0, fn=None): """ return a new sequence derived from input sequence , but with elements removed from the front and elements removed from the end. >>> trim([1, 2, 3, 4, 5], head=2) [3, 4, 5] >>> trim([1, 2, 3, 4, 5], tail=2) [1, 2, 3] >>> trim([1, 2, 3, 4, 5], head=2, tail=2) [3] >>> trim('progress', head=2, tail=2) 'ogre' """ if head > 0 or tail > 0: if fn is None: if isinstance(seq, basestring): fn = join elif isinstance(seq, tuple): fn = tuple seq = list(seq) if head > 0: del seq[:head] if tail > 0: del seq[-tail:] return (fn(seq) if fn else seq) def _partitions(seq, n): """ partition a sequence of distinct elements into subsequences of length . should be sequenceable type (tuple, list, str). should be a factor of the size of the sequence. >>> list(_partitions((1, 2, 3, 4), 2)) [((1, 2), (3, 4)), ((1, 3), (2, 4)), ((1, 4), (2, 3))] """ if not (len(seq) > n): yield (seq,) else: for x in itertools.combinations(seq[1:], n - 1): p = (seq[0],) + tuple(x) for ps in _partitions(diff(seq[1:], x), n): yield (p,) + ps def ipartitions(seq, n): """ partition a sequence by index. >>> list(ipartitions((1, 0, 1, 0), 2)) [((1, 0), (1, 0)), ((1, 1), (0, 0)), ((1, 0), (0, 1))] """ for p in _partitions(tuple(xrange(len(seq))), n): yield tuple(tuple(seq[i] for i in x) for x in p) def partitions(seq, n, pad=0, value=None, distinct=None): """ partition a sequence into subsequences of length . if is true then the sequence will be padded (using ) until its length is a integer multiple of . if sequence contains distinct elements then can be set to True, if it is not set then will be examined for repeated elements. >>> list(partitions((1, 2, 3, 4), 2)) [((1, 2), (3, 4)), ((1, 3), (2, 4)), ((1, 4), (2, 3))] """ if not isinstance(seq, (tuple, list, str)): seq = tuple(seq) (d, r) = divmod(len(seq), n) if r != 0: if not pad: raise ValueError("invalid sequence length {l} for {n}-tuples".format(l=len(seq), n=n)) seq = tuple(seq) + (value,) * (n - r) if d == 0 or (d == 1 and r == 0): yield (seq,) else: if distinct is None: distinct = is_pairwise_distinct(*seq) fn = (_partitions if distinct else ipartitions) # or in Python 3: [[ yield from fn(seq, n) ]] for z in fn(seq, n): yield z # see: [ https://enigmaticcode.wordpress.com/2017/05/17/tantalizer-482-lapses-from-grace/#comment-7169 ] # choose: choose values from satisfying in turn # distinct - true if values must be distinct # s - initial sequence (that supports 'copy()' and 'append()') def choose(vs, fns, s=None, distinct=0): """ choose values from satisfying in turn. if all values are acceptable then a value of None can be passed in . set 'distinct' if all values should be distinct. >>> list(choose([1, 2, 3], [None, (lambda a, b: abs(a - b) == 1), (lambda a, b, c: abs(b - c) == 1)])) [[1, 2, 1], [1, 2, 3], [2, 1, 2], [2, 3, 2], [3, 2, 1], [3, 2, 3]] """ if s is None: s = list() # are we done? if not fns: yield s else: # choose the next value fn = fns[0] for v in vs: if (not distinct) or v not in s: s_ = list(s) s_.append(v) if fn is None or fn(*s_): # choose the rest [[Python 3: yield from ...]] for z in choose(vs, fns[1:], s_, distinct): yield z def first(s, count=1, skip=0, fn=list): """ return the first items in iterator (skipping the initial items) as a list (or other object specified by ). can be a callable object, in which case items are collected from while returns a true value when it is passed each item (after skipping the first items). can also be a callable, in which case items are skipped while returns a true value when it is passed each item. this would be a way to find the first 10 primes: >>> first((n for n in irange(1, inf) if is_prime(n)), count=10) [2, 3, 5, 7, 11, 13, 17, 19, 23, 29] >>> first(p for p in primes if p % 17 == 1) [103] this finds squares less than 200 >>> first((sq(x) for x in irange(0, inf)), count=(lambda x: x < 200)) [0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196] """ if callable(count): if skip == 0: r = itertools.takewhile(count, s) elif callable(skip): r = itertools.takewhile(count, itertools.dropwhile(skip, s)) else: r = itertools.takewhile(count, itertools.islice(s, skip, None)) elif count == inf: r = s else: r = itertools.islice(s, skip, skip + count) return (r if fn is None else fn(r)) # return the single value if s contains only a single value (else None) # NOTE: similar to the Python expression : [[ [x] = s ]] def singleton(s, skip=0, default=None): """ if the container contains only a single value return it, otherwise return None (or the parameter) >>> singleton([], default=0) 0 >>> singleton({1}, default=0) 1 >>> singleton([1, 2, 3], default=0) 0 """ r = first(s, 2, skip) return (r[0] if len(r) == 1 else default) def repeat(fn, v=0, k=inf): """ generate repeated applications of function to value . the initial value is returned first, followed by the result of repeatedly applying the specified function to the previous value. if a limit is specified then the function will be applied the specified number of times, so (k + 1) values will be returned (corresponding to the application of the function 0 .. k times). >>> list(repeat((lambda x: x + 1), 0, 10)) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] """ i = 0 while True: yield v if i == k: break i += 1 v = fn(v) def uniq(seq, fn=None, verbose=0): """ generate unique values from (maintaining order). i.e. repeated values are suppressed. >>> list(uniq([5, 7, 0, 0, 5])) [5, 7, 0] >>> join(uniq('mississippi')) 'misp' >>> list(uniq(irange(1, 9), fn=(lambda x: x // 3))) [1, 3, 6, 9] """ seen = set() for x in seq: r = (x if fn is None else fn(x)) if r not in seen: yield x seen.add(r) if verbose: printf("[uniq: found {n} unique items]", n=len(seen)) def uniq1(seq, fn=None): """ collapse repeated consecutive values (according to ) in down to single values. i.e. repeated _consecutive_ values are suppressed. >>> list(uniq1((1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5))) [1, 2, 3, 4, 5] >>> join(uniq1('mississippi')) 'misisipi' >>> join(uniq1('bookkeeper')) 'bokeper' """ for vs in clump(seq, fn=fn): yield vs[0] # root: calculate the (positive) nth root of a (positive) number # we use math.pow rather than **/pow() to avoid generating complex numbers root = lambda x, n: (x if not x else math.pow(x, 1.0 / n)) def cbrt(x): """ return the cube root of a number (as a float). >>> cbrt(27.0) 3.0 >>> cbrt(-27.0) -3.0 """ r = root(abs(x), 3.0) return (-r if x < 0 else r) # cb = lambda x: x**3 def cb(x): "cb(x) = x**3"; return x**3 # for large numbers try Primes.prime_factor(n, mr=100), or sympy.ntheory.factorint(n) # basis = [2, 3, 5] _prime_factor_ds = (1, 2, 2, 4, 2, 4, 2, 4, 6, 2, 6) # deltas _prime_factor_js = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 3) # next index def prime_factor(n): """ generate (, ) pairs in the prime factorisation of positive integer . no pairs are returned for 1 (or for non-positive integers). for numbers with large prime factors it will take a long time to find them. >>> list(prime_factor(60)) [(2, 2), (3, 1), (5, 1)] >>> list(prime_factor(factorial(12))) [(2, 10), (3, 5), (5, 2), (7, 1), (11, 1)] """ if n > 1: ds = _prime_factor_ds js = _prime_factor_js x = 2 j = 0 while True: e = 0 while True: (d, r) = divmod(n, x) if r > 0: break e += 1 n = d if e > 0: yield (x, e) x += ds[j] if x * x > n: break j = js[j] # anything left is prime if n > 1: yield (n, 1) # maybe should be called factors() or factorise() def factor(n, fn=prime_factor): """ return a list of the prime factors of positive integer . for integers less than 1, None is returned. The parameter is used to generate the prime factors of the number. (Defaults to using prime_factor()). >>> factor(101) [101] >>> factor(1001) [7, 11, 13] >>> factor(12) [2, 2, 3] >>> factor(125) [5, 5, 5] """ if n < 1: return None factors = [] for (p, e) in fn(n): factors.extend([p] * e) return factors # divsors (for non-negative integers) are based on: # # (a, b) is a divisor pair of n iff: a, b in [0, n] and a.b = n # # the set of divisors is the set of numbers that appear in the divisor pairs # # and these are returned in order, # # so: # # divisor pairs 6 = (1, 6) (2, 3); divisors 6 = (1, 2, 3, 6) # divisor pairs 4 = (1, 4) (2, 2); divisors 4 = (1, 2, 4) # divisor pairs 2 = (1, 2); divisors 2 = (1, 2) # divisor pairs 1 = (1, 1); divisors 1 = (1) # divisor pairs 0 = (0, 0); divisors 0 = (0) # is a divisor of ? def is_divisor(n, x, proper=0): """ determine if is a divisor of (both are non-negative integers). if 'proper' is set then the divisor must be smaller than . >>> is_divisor(42, 7) True >>> is_divisor(43, 7) False >>> is_divisor(7, 7) True >>> is_divisor(7, 7, proper=1) False >>> is_divisor(1, 0) False >>> is_divisor(0, 0) True """ if x < 0 or x > n: return False if n == 0: return (not proper) and (x == 0) return (x > 0 and n % x == 0 and ((not proper) or x < n)) # you can use the following to look for multiples but note argument order is opposite of div() # and this will only work for multiples from 1 upwards (or from 2 if 'proper' is set) #is_multiple = lambda n, x, proper=0: is_divisor(x, n, proper=proper) def divisor_pairs(n): """ generate divisors (a, b) of positive integer n, such that a <= b and a * b = n. the pairs are generated in order of increasing . if you only want a few small divisors, this routine is OK, otherwise you are probably better using divisors_pairs(). >>> list(divisor_pairs(36)) [(1, 36), (2, 18), (3, 12), (4, 9), (6, 6)] >>> list(divisor_pairs(101)) [(1, 101)] """ if n == 0: yield (0, 0) return a = 1 while True: (b, r) = divmod(n, a) if a > b: break if r == 0: yield (a, b) a += 1 def divisor(n): """ generate divisors of positive integer in numerical order. >>> list(divisor(36)) [1, 2, 3, 4, 6, 9, 12, 18, 36] >>> list(divisor(101)) [1, 101] """ bs = list() for (a, b) in divisor_pairs(n): yield a if a < b: bs.insert(0, b) else: break for b in bs: yield b def multiples(ps, k=1): """ given a list of (, ) pairs, return all numbers that can be formed by multiplying together the s, with each occurring up to * times. the multiples are returned as a sorted list the practical upshot of this is that the divisors of a number can be found using the expression: multiples(prime_factor(x)) >>> multiples(prime_factor(180)) [1, 2, 3, 4, 5, 6, 9, 10, 12, 15, 18, 20, 30, 36, 45, 60, 90, 180] """ s = [1] for (m, n) in ps: if k > 1: n *= k t = list() p = m for _ in xrange(n): t.extend(x * p for x in s) p *= m s.extend(t) s.sort() return s def divisors(n, k=1, fn=prime_factor): """ return the divisors of positive integer pow(, ) as a sorted list. >>> divisors(36) [1, 2, 3, 4, 6, 9, 12, 18, 36] >>> divisors(101) [1, 101] """ if n == 0 and k > 0: return [0] return multiples(fn(n), k=k) def divisors_pairs(n, k=1, fn=prime_factor, every=0): """ generate divisors pairs (a, b) with a <= b, such that a * b = pow(n, k). pairs are generated in order, by determining the factors of n. this is probably faster than divisor_pairs() if you want all divisors. if the 'every' parameter is set, then pairs with a > b are also generated. """ if n == 0 and k > 0: yield (0, 0) return nk = (n**k if k != 1 else n) for a in divisors(n, k=k, fn=fn): b = nk // a if a > b and (not every): break yield (a, b) def divisors_tuples(n, k, s=()): """ find ordered -tuples that multiply to give . >>> list(divisors_tuples(1335, 3)) [(1, 1, 1335), (1, 3, 445), (1, 5, 267), (1, 15, 89), (3, 5, 89)] """ if k == 1: if not (s and n < s[-1]): yield s + (n,) else: for (a, b) in divisors_pairs(n): if not (s and a < s[-1]): for z in divisors_tuples(b, k - 1, s + (a,)): yield z def is_prime(n): # type: (int) -> bool """ return True if the positive integer is prime. Note: for numbers up to 2**64 is_prime_mr() is a fast, accurate prime test. (And for larger numbers it is probabilistically accurate). >>> is_prime(101) True >>> is_prime(1001) False """ if n < 2: return False # 0, 1 -> F if n < 4: return True # 2, 3 -> T r = n % 6 if r != 1 and r != 5: return False # (n % 6) != (1, 5) -> F for (p, e) in prime_factor(n): return p == n return False prime = is_prime # Miller-Rabin primality test (originally suggested by Brian Gladman) import random def _is_composite(a, d, n, s): if a == 0: return 0 x = pow(a, d, n) if x == 1: return 0 for _ in xrange(s): if x == n - 1: return 0 x = (x * x) % n # definitely composite return 1 def is_prime_mr(n, r=0): """ Miller-Rabin primality test for . is the number of random extra rounds performed for large numbers return value: 0 = the number is definitely not prime (definitely composite for n > 1) 1 = the number is probably prime 2 = the number is definitely prime for numbers less than 2**64, the prime test is completely accurate, and deterministic, the extra rounds are not performed. for larger numbers additional rounds are performed, and if the number cannot be found to be composite a value of 1 (probably prime) is returned. confidence can be increased by using more additional rounds. >>> is_prime_mr(288230376151711813) 2 >>> is_prime_mr(316912650057057350374175801351) 1 >>> is_prime_mr(332306998946228968225951765070086171) 0 """ # 0, 1 = not prime if n < 2: return 0 # 2, 3 = definitely prime if n < 4: return 2 # all other primes have a residue mod 6 of 1 or 5 x = n % 6 if x != 1 and x != 5: return 0 # compute 2^s.d = n - 1 d = n - 1 s = (d & -d).bit_length() - 1 d >>= s # bases from: https://miller-rabin.appspot.com/ # we use 3 sets of bases: # 1 base = [9345883071009581737] is completely accurate for n < 341531 (about 2^18) # 2 bases = [336781006125, 9639812373923155] (2 bases) is completely accurate for n < 1050535501 (about 2^30) # 7 bases = [2, 325, 9375, 28178, 450775, 9780504, 1795265022] is completely accurate for n < 2^64 # 1 base is completely accurate for n < 341531 if n < 341531: return (0 if _is_composite(9345883071009581737 % n, d, n, s) else 2) # 2 bases are completely accurate for n < 1050535501 if n < 1050535501: return (0 if _is_composite(336781006125 % n, d, n, s) or _is_composite(9639812373923155 % n, d, n, s) else 2) # test remaining numbers with the 7 base set if any(_is_composite(a % n, d, n, s) for a in (2, 325, 9375, 28178, 450775, 9780504, 1795265022)): # definitely composite return 0 # the 7 base set is completely accurate for n < 2^64: if n < 0x10000000000000000: # definitely prime return 2 # for larger numbers run further prime tests as specified if r > 0 and any(_is_composite(random.randrange(2, n - 1), d, n, s) for _ in range(r)): # definitely composite return 0 # otherwise, probably prime return 1 # find prime factors using Pollard's Rho method # [see: https://programmingpraxis.files.wordpress.com/2012/09/primenumbers.pdf ] def _rho_factor(n, mrr=0): if n % 2 == 0: return 2 c = 1 while True: (t, h, d) = (2, 2, 1) while d == 1: t = (t * t + c) % n h = (h * h + c) % n h = (h * h + c) % n d = gcd(h - t, n) if d == n: pass elif is_prime_mr(d, mrr): return d else: n = d c += 1 # NOTE: factors are not neccessarily returned in order def prime_factor_rho(n, mrr=0): """ generate (, ) pairs in the prime factorisation of positive integer . note that factors are not necessarily returned in numerical order. is the number of additional rounds performed in the is_prime_mr() test for prime factors. >> sorted(prime_factor_rho(factorial(23) + 1)) [(47, 2), (79, 1), (148139754736864591, 1)] """ while n > 1: # check for prime if is_prime_mr(n, mrr): yield (n, 1) break # find a factor p = _rho_factor(n) (e, n) = drop_factors(n, p) yield (p, e) def prime_factor_h(n, ps=None, end=None, nf=0, mr=0, mrr=0): """ find prime factors using various heuristics ps = can be a prime sieve to check (end = upper limit on primes) nf = number of tolerated failures (after which we switch to heuristics) mr = enable Pollard Rho/Miller Rabin (mrr = number of Miller-Rabin rounds) Primes found using the sieve will be generated in numerical order. Primes found heuristically may not be generated in order. Depending on the arguments the factorisation may be incomplete (e.g. if a sieve is specified and mr=0, large factors outside the sieve may not be found). The following example uses the prime sieve for factors up to 1000, and then probabalistic tests to find the rest: >> list(prime_factor_h(factorial(18) + 1, ps=primes, end=1000, mr=1)) [(19, 1), (23, 1), (29, 1), (61, 1), (67, 1), (123610951, 1)] """ assert not (ps is None and mr == 0) # otherwise we won't find anything f = 0 # number of failed primes for n psi = (None if ps is None else ps.generate(end=end)) pmax = 0 # first use primes from the sieve while psi is not None and n > 1: # is this a prime in the sieve? if n < ps.max: if f == 0 and ps.is_prime(n): yield (n, 1) return # try the next prime in the sieve try: p = next(psi) except StopIteration: # the sieve is exhausted, so n is a product of primes larger than the sieve pmax = end or ps.max if n < pmax * pmax: # n must be a prime yield (n, 1) return else: # move to probabalistic tests psi = None else: # check prime p (e, n) = drop_factors(n, p) if e > 0: yield (p, e) f = 0 else: f += 1 if f == nf: # move to probabalistic tests pmax = end or ps.max psi = None # now try heuristic tests on what is left if mr and n > 1: m = 1 # multiplicity of factors # check to see if the number is an exact power k = 2 while True: r = iroot(n, k) if pow(r, k) == n: m *= k n = r elif k == 2: k = 3 else: k += 2 # limit the search if k > 20 or r < pmax: break # n could now be in the sieve if ps is not None and n < pmax * pmax: yield (n, m) return # look for probabalistic factors (not necessarily in order) for (p, e) in prime_factor_rho(n, mrr=mrr): yield (p, e * m) def tau(n, fn=prime_factor): """ count the number of divisors of a positive integer . tau(n) = len(divisors(n)) [but faster] >>> tau(factorial(12)) 792 """ return multiply(e + 1 for (_, e) in fn(n)) def is_square_free(n, fn=prime_factor): """ a positive integer is "square free" if it is not divisibly by a perfect square greater than 1. >>> is_square_free(8596) False >>> is_square_free(8970) True """ return n > 0 and all(e == 1 for (_, e) in fn(n)) def mobius(n, fn=prime_factor): """ return the Mobius value for positive integer . mobius(n) = 1; if n is square free and has an even number of prime factors mobius(n) = -1; if n is square free and has an odd number of prime factors mobius(n) = 0; if n is not square free """ if n < 1: return None r = 1 for (p, e) in fn(n): if e > 1: return 0 r = -r return r def farey(n, ends=0): """ generate the Farey sequence F(n) - the sequence of coprime pairs (a, b) where 0 < a < b <= n. pairs are generated in numerical order when considered as fractions a/b. the pairs (0, 1) and (1, 1) usually present at the start and end of the sequence are not generated by this function, unless 'ends' is set to True. >>> list(p for p in farey(20) if sum(p) == 20) [(1, 19), (3, 17), (7, 13), (9, 11)] """ if ends: yield (0, 1) (a, b, c, d) = (0, 1, 1, n) while d > 1: k = (n + b) // d (a, b, c, d) = (c, d, k * c - a, k * d - b) yield (a, b) if ends: yield (1, 1) def coprime_pairs(n=None, order=0): """ generate coprime pairs (a, b) with 0 < a < b <= n. the list is complete and no element appears more than once. if n is not specified then pairs will be generated indefinitely. if n is specified then farey() can be used instead to generate coprime pairs in numerical order (when considered as fractions). if order=1 is specified then the pairs will be produced in order. >>> sorted(p for p in coprime_pairs(20) if sum(p) == 20) [(1, 19), (3, 17), (7, 13), (9, 11)] >>> list(coprime_pairs(6, order=1)) [(1, 2), (1, 3), (2, 3), (1, 4), (3, 4), (1, 5), (2, 5), (3, 5), (4, 5), (1, 6), (5, 6)] """ fn = ((lambda p: p[0] <= n) if n else true) if order: # use a heap to order the pairs from heapq import (heapify, heappush, heappop) ps = list() heapify(ps) _push = heappush _pop = heappop else: # just use a list ps = list() _push = lambda ps, p: ps.append(p) _pop = lambda ps: ps.pop(0) for p in ((2, 1), (3, 1)): if fn(p): _push(ps, p) while ps: (b, a) = _pop(ps) yield (a, b) for p in ((2 * b - a, b), (2 * a + b, a), (2 * b + a, b)): if fn(p): _push(ps, p) # Pythagorean Triples: # see: https://en.wikipedia.org/wiki/Formulas_for_generating_Pythagorean_triples # generate primitive pythagorean triples (x, y, z) with hypotenuse not exceeding Z # if Z is None, then triples will be generated indefinitely # if order is true, then triples will be returned in order def _pythagorean_primitive(Z=None, order=0): fn = (true if Z is None else le(Z)) if order: # use a heap from heapq import (heapify, heappush, heappop) ts = list() heapify(ts) _push = heappush _pop = heappop else: # just use a list ts = list() _push = lambda s, t: s.append(t) _pop = lambda s: s.pop(0) # initial triple if fn(5): _push(ts, (5, 4, 3)) while ts: (c, b, a) = _pop(ts) yield (a, b, c) # my original formulation (using only addition/subtraction) (a2, b2, c2) = (a + a, b + b, c + c) c3 = c2 + c for (z, y, x) in ( (c3 - b2 + a2, c2 - b + a2, c2 - b2 + a), (c3 + b2 - a2, c2 + b - a2, c2 + b2 - a), (c3 + b2 + a2, c2 + b + a2, c2 + b2 + a), ): if fn(z): _push(ts, ((z, x, y) if y < x else (z, y, x))) ## alternatively: Brian's (more compact, but slower) formulation #t = 2 * (a + b + c) #(u, v, w) = (t - 4 * b, t, t - 4 * a) #for (z, y, x) in ((u + c, u + b, u - a), (v + c, v - b, v - a), (w + c, w - b, w + a)): # if fn(z): _push(ts, ((z, x, y) if y < x else (z, y, x))) # generate pythagorean triples (x, y, z) with hypotenuse not exceeding Z def _pythagorean_all(Z, order=0): if order: # use a heap to save the multiples from heapq import (heapify, heappush, heappop) ms = list() heapify(ms) for (x, y, z) in _pythagorean_primitive(Z, order=1): # return any saved multiples less than (x, y, z) while ms and ms[0] < (z, y, x): yield heappop(ms)[::-1] # return (x, y, z) yield (x, y, z) # add in any new multiples for k in irange(2, Z // z): heappush(ms, (k * z, k * y, k * x)) # return any remaining multiples while ms: yield heappop(ms)[::-1] else: # return the multiples with the primitives for (x, y, z) in _pythagorean_primitive(Z, order=0): yield (x, y, z) for k in irange(2, Z // z): yield (k * x, k * y, k * z) # generate pythagorean triples # n - specifies the maximum hypotenuse allowed # primitive - if set only primitive triples are generated # order - if set triples are generated in order # if primitive is false, then a value for n must be specified def pythagorean_triples(n=None, primitive=0, order=0): """ generate pythagorean triples (x, y, z) where x < y < z and x^2 + y^2 = z^2. n - maximum allowed hypotenuse (z) primitive - if set only primitive triples are generated order - if set triples are generated in order order is by shortest z, then shortest y, then shortest x (i.e. reverse lexicographic) if 'primitive' is set, then n can be None, and primitive triples will be generated indefinitely (although it will eventually run out of memory) >>> list(pythagorean_triples(20, primitive=0, order=1)) [(3, 4, 5), (6, 8, 10), (5, 12, 13), (9, 12, 15), (8, 15, 17), (12, 16, 20)] >>> list(pythagorean_triples(20, primitive=1, order=1)) [(3, 4, 5), (5, 12, 13), (8, 15, 17)] >>> icount(pythagorean_triples(10000, primitive=1)) 1593 >>> icount(pythagorean_triples(10000, primitive=0)) 12471 """ if primitive: # primitive only triples return _pythagorean_primitive(n, order) else: # include non-primitive assert n is not None return _pythagorean_all(n, order) def fib(*s, **kw): """ generate Fibonacci type sequences (or other recurrence relations) The initial k terms are provided as sequence s, subsequent terms are calculated as a function of the preceeding k terms. The default function being 'sum', but a different function can be specified using the 'fn' parameter (which should be a function that takes a sequence of k terms and computes the appropriate value). Standard Fibonacci numbers (OEIS A000045): >>> first(fib(0, 1), 10) [0, 1, 1, 2, 3, 5, 8, 13, 21, 34] Lucas numbers (OEIS A000032): >>> first(fib(2, 1), 10) [2, 1, 3, 4, 7, 11, 18, 29, 47, 76] Tribonacci numbers (OEIS A001590): >>> first(fib(0, 1, 0), 10) [0, 1, 0, 1, 2, 3, 6, 11, 20, 37] Powers of 2 (using addition): >>> first(fib(1, fn=unpack(lambda x: x + x)), 10) [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] """ fn = kw.get('fn', sum) s = list(s) while True: s.append(fn(s)) yield s.pop(0) # if we don't overflow floats (happens around 2^53) this works... # def is_power(n, m): # i = int(n**(1.0 / m) + 0.5) # return (i if i**m == n else None) # but here we use a binary search, which should work on arbitrary large integers # # NOTE: that this will return 0 if n = 0 and None if n is not a perfect k-th power, # so [[ power(n, k) ]] will evaluate to True only for positive n # if you want to allow n to be 0 you should check: [[ power(n, k) is not None ]] def iroot(n, k): """ compute the largest integer x such that pow(x, k) <= n. i.e. x is the integer k-th root of n. it is the exact root if: pow(x, k) == n (which is what is_power() does) """ # binary search if n < 0 or k < 1: return if n >> k == 0: return int(n > 0) a = 1 << ((n.bit_length() - 1) // k) b = a << 1 #assert (a**k <= n and b**k > n) # if this assertion fails we need: #while not (b**k > n): (a, b) = (b, b << 1) while b - a > 1: r = (a + b) // 2 x = r**k if x < n: a = r elif x > n: b = r else: return r return a def is_power(n, k): """ check positive integer is a perfect th power of some integer. if is a perfect th power, returns the integer th root. if is not a perfect th power, returns None. >>> is_power(49, 2) 7 >>> is_power(49, 3) is not None False >>> is_power(0, 2) 0 >>> n = (2**60 + 1) >>> (is_power(n**2, 2) is not None, is_power(n**2 + 1, 2) is not None) (True, False) >>> (is_power(n**3, 3) is not None, is_power(n**3 + 1, 3) is not None) (True, False) """ r = iroot(n, k) if r is None: return None return (r if r**k == n else None) def sqrt(a, b=None): """ the (real) square root of a / b (or just a if b is None) >>> sqrt(9) 3.0 >>> sqrt(9, 4) 1.5 """ # / is operator.truediv() here return math.sqrt(a if b is None else a / b) # sq = lambda x: x * x def sq(x): "sq(x) = x**2"; return x * x def sumsq(xs): "sumsq(xs) = sum(sq(x) for x in xs)"; return sum(x * x for x in xs) # calculate intf(sqrt(n)) @static(impl=getattr(math, 'isqrt', None)) def isqrt(n): # type: (int) -> int | NoneType """ calculate intf(sqrt(n)), for integers n. See also: math.isqrt (Python 3.8+), gmpy2.isqrt(). >>> isqrt(9) 3 >>> isqrt(15) 3 >>> isqrt(16) 4 >>> isqrt(17) 4 """ if n < 0: return None if n < 4: return int(n > 0) if isqrt.impl: return isqrt.impl(n) # use math.isqrt() if available # use the math.isqrt algorithm c = (n.bit_length() - 1) >> 1 a = 1 d = 0 s = c.bit_length() while s: s -= 1 e = d d = c >> s a = (a << d - e - 1) + (n >> (c << 1) - e - d + 1) // a return a - (a * a > n) # square root floor and ceiling functions sqrtf = isqrt sqrtc = lambda x: (isqrt(x) if x < 1 else 1 + isqrt(x - 1)) # it would be more Pythonic to encapsulate is_square in a class with the initialisation # in __init__, and the actual call in __call__, and then instantiate an object to be # the is_square() function (i.e. [[ is_square = _is_square_class(80) ]]), but it is # more efficient (and perhaps more readable) to just use normal variables, although # if you're using PyPy the class based version is just as fast (if not slightly faster) # experimentally mod = 80, 48, 72, 32 are good values (24, 16 also work OK) @static(mod=720, residues=None, cache_enabled=0, cache=dict()) def is_square(n): # type: (int) -> int | NoneType """ check integer is a perfect square. if is a perfect square, returns the integer square root. if is not a perfect square, returns None. results can be cached by setting: is_square.cache_enabled = 1 >>> is_square(49) 7 >>> is_square(50) is not None False >>> is_square(0) 0 """ if n < 0: return None if n < 2: return n # early rejection: check mod against a precomputed cache # e.g. mod 80 = 0, 1, 4, 9, 16, 20, 25, 36, 41, 49, 64, 65 (rejects 88% of numbers) # mod 720 (= factorial(6)) rejects 93% of candidates if not is_square.residues: is_square.residues = set((i * i) % is_square.mod for i in xrange(is_square.mod)) if (n % is_square.mod) not in is_square.residues: return None # otherwise use isqrt and check the result z = is_square.cache.get(n) if z is None: r = isqrt(n) z = (r if r * r == n else None) if is_square.cache_enabled: is_square.cache[n] = z return z # is the square of a rational number? def is_square_q(n, F=None): if F is None: F = Rational() n = F(n) p = is_square(n.numerator) if p is None: return None q = is_square(n.denominator) if q is None: return None return F(p, q) def sum_of_squares(n, k=2, min_v=0, sep=0, ss=[]): """ return ordered k-sequences of non-negative integers (a, b, ...) such that: n = a**2 + b**2 + ... min_v - specifies the minimum allowable value in the returned sequences sep - specified the minimum separation between values >>> list(sum_of_squares(50, 2)) [[1, 7], [5, 5]] >>> list(sum_of_squares(50, 2, sep=1)) [[1, 7]] >>> list(sum_of_squares(637, 3)) [[0, 14, 21], [3, 12, 22], [5, 6, 24], [12, 13, 18]] """ if k == 1: r = is_square(n) if not (r is None or r < min_v): yield ss + [r] elif k == 2: i = isqrt(n) j = 0 while not (i < j): r = compare(i * i + j * j, n) if r == 0 and not (j < min_v or i - j < sep): yield (ss + [j, i] if ss else [j, i]) if r != -1: i -= 1 if r != 1: j += 1 else: for x in irange(min_v, inf): x2 = x * x if x2 * (k - 1) > n: break for z in sum_of_squares(n - x2, k - 1, x + sep, sep, ss + [x]): yield z # generate powers from a range def powers(a, b, k=2, step=1, fn=None): """ generate powers pow(n, k) for n in irange(a, b) >>> list(powers(1, 10)) [1, 4, 9, 16, 25, 36, 49, 64, 81, 100] >>> list(powers(1, 10, 3)) [1, 8, 27, 64, 125, 216, 343, 512, 729, 1000] """ for n in irange(a, b, step=step): x = n**k yield (x if fn is None else fn(x)) # compose functions in order (forward functional composition, "and then") # so: fcompose(f, g, h)(x) == h(g(f(x))) def fcompose(f, *gs): """ forward functional composition ("and then") fcompose(f, g, h)(x) == h(g(f(x))) >>> fcompose(is_square, is_not_none)(49) True >>> fcompose(is_square, is_not_none)(50) False """ # special case for 1 or 2 functions n = len(gs) if n == 0: return f if n == 1: g = gs[0] return (lambda *args, **kw: g(f(*args, **kw))) # general case def fn(*args, **kw): r = f(*args, **kw) for g in gs: r = g(r) return r return fn # compose functions in reverse order (reverse functional composition, "after") # so: rcompose(f, g, h)(x) = f(g(h(x))) def rcompose(*fns): """ reverse functional composition rcompose(f, g, h)(x) == f(g(h(x))) """ return fcompose(*(reversed(fns))) is_not_none = (lambda x: x is not None) is_square_p = (lambda x: is_square(x) is not None) # = fcompose(is_square, is_not_none) # 819 rejects 95% (other good values: 63 (86%), 117 (87%), 189 (89%), 351 (90%), 504 (91%), 819 (95%)) @static(mod=819, residues=None, cache_enabled=0, cache=dict()) def is_cube(n): """ check positive integer is a perfect cube. to check for positive/negative values use: is_cube_z(). results can be cached by setting: is_cube.cache_enabled = 1 >>> is_cube(27) 3 >>> is_cube(49) is not None False >>> is_cube(0) 0 """ if n < 0: return None if n < 2: return n if not is_cube.residues: is_cube.residues = set((i * i * i) % is_cube.mod for i in xrange(is_cube.mod)) if (n % is_cube.mod) not in is_cube.residues: return None z = is_cube.cache.get(n) if z is None: z = is_power(n, 3) if is_cube.cache_enabled: is_cube.cache[n] = z return z is_cube_p = (lambda x: is_cube(x) is not None) # = fcompose(is_cube, is_not_none) def is_cube_z(n): """ check integer is a perfect cube. >>> is_cube_z(27) 3 >>> is_cube_z(-27) -3 >>> is_cube_z(0) 0 """ if n < 0: r = is_cube(-n) return (r if r is None else -r) else: return is_cube(n) # keep the old names as aliases power = is_power cube = is_cube square = is_square def drop_factors(n, k): """ remove factors of from . return (i, m) where n = (m)(k^i) such that m is not divisible by k """ i = 0 while n > 1: (d, r) = divmod(n, k) if r != 0: break i += 1 n = d return (i, n) def is_power_of(n, k): """ check is a power of . returns such that pow(k, m) = n or None. >>> is_power_of(128, 2) 7 >>> is_power_of(1, 2) 0 >>> is_power_of(0, 2) is None True >>> is_power_of(0, 0) 1 """ if n == 0: return (1 if k == 0 else None) if n == 1: return 0 if k < 2: return None (i, m) = drop_factors(n, k) return (i if m == 1 else None) def tri(n): """ tri(n) is the nth triangular number. tri(n) = n * (n + 1) / 2. Note: trif() is available for float arguments. >>> tri(1) 1 >>> tri(100) 5050 """ return n * (n + 1) // 2 T = tri # triangular numbers as floats trif = lambda x: 0.5 * x * (x + 1) def trirt(x): """ return the triangular root of (as a float) >>> trirt(5050) 100.0 >>> round(trirt(2), 8) 1.56155281 """ return 0.5 * (math.sqrt(8 * x + 1) - 1.0) def is_triangular(n): """ check positive integer is a triangular number. if is a triangular number, returns integer such that T(k) == n. if is not a triangular number, returns None. >>> is_triangular(5050) 100 >>> is_triangular(49) is not None False """ if n % 9 not in {0, 1, 3, 6}: return x = is_square(8 * n + 1) return (None if x is None else x // 2) def digrt(n, base=10): """ return the digital root of positive integer . >>> digrt(123456789) 9 >>> digrt(sum([1, 2, 3, 4, 5, 6, 7, 8, 9])) 9 >>> digrt(factorial(100)) 9 """ return (0 if n == 0 else int(1 + (n - 1) % (base - 1))) def repdigit(n, d=1, base=10): """ return a number consisting of the digit repeated times, in base default digit is d=1 (to return repunits) default base is base=10 >>> repdigit(6) 111111 >>> repdigit(6, 7) 777777 >>> repdigit(6, 7, base=16) 7829367 >>> repdigit(6, 7, base=16) == 0x777777 True """ assert 0 <= d < base return d * ((base**n) - 1) // (base - 1) # Python 3.6: ...(*vs, root=math.sqrt) def hypot(*vs, **kw): """ return hypotenuse of a right angled triangle with shorter sides and . hypot(a, b) = sqrt(a^2 + b^2) multiple arguments can be specified to return Euclidean distance in higher dimensions. a keyword argument of 'root' may be specified to provide the function used to calculate the root of the sum of the squares. See also: math.hypot() (Python 3.8+). >>> hypot(3, 4) 5.0 >>> hypot(3, 4, 12) 13.0 >>> hypot(3, 4, root=is_square) 5 """ root = kw.get('root', math.sqrt) return root(sum(v * v for v in vs)) # alias for: hypot(..., root=is_square) ihypot = lambda *vs: hypot(*vs, root=is_square) # return roots of the form n/d in the appropriate domain def _roots(domain, F, *nds): for (n, d) in nds: if domain in "CF": yield F(n / d) elif domain in "QZ": x = F(n, d) if domain == "Q": yield x elif domain == "Z": x = as_int(x, default=None) if x is not None: yield x # find roots of a quadratic equation # domain = "Z" (integer), "Q" (rational), "F" (float), "C" (complex float) def quadratic(a, b, c, domain="Q", F=None): """ find roots of the equation: a.x^2 + b.x + c = 0 in the specified domain: "Z" finds integer solutions "Q" finds rational solutions "F" finds float solutions "C" finds complex solutions >>> sorted(quadratic(1, 1, -6, domain="Z")) [-3, 2] """ assert domain in "CFQZ", sprintf("quadratic: invalid domain '{domain}'") if a == 0: # linear equation assert b != 0 if F is None: F = (Rational() if domain in 'QZ' else (complex if domain == 'C' else float)) return _roots(domain, F, (-c, b)) # discriminant D = b * b - 4 * a * c if D < 0 and domain != "C": return _roots(domain, None) if domain in "CF": F = (complex if domain == 'C' else float) d = -2 * a if D == 0: return _roots(domain, F, (b, d)) else: r = D**0.5 return _roots(domain, F, (b + r, d), (b - r, d)) elif domain in "QZ": if F is None: F = Rational() r = is_square_q(D, F=F) if r is not None: d = -2 * a if r == 0: return _roots(domain, F, (b, d)) else: return _roots(domain, F, (b + r, d), (b - r, d)) return _roots(domain, None) def intf(x): """ floor conversion (largest integer not greater than x) of float to int. >>> intf(1.0) 1 >>> intf(1.5) 1 >>> intf(-1.5) -2 """ r = int(x) return (r - 1 if x < r else r) def intc(x): """ ceiling conversion (smallest integer not less than x) of float to int. >>> intc(1.0) 1 >>> intc(1.5) 2 >>> intc(-1.5) -1 """ r = int(x) return (r + 1 if x > r else r) def intr(x): """ round to nearest integer. if x is exactly between two integers (i.e. x = n + 0.5) then the answer is the integer further away from 0. >>> intr(0.0) 0 >>> intr(2.5) 3 >>> intr(-2.5) -3 """ if x < 0: x = -x return -int((x + x + 1) // 2) else: return int((x + x + 1) // 2) def divf(a, b): """ floor division. (similar to Python's a // b). always returns an int. >>> divf(100, 10) 10 >>> divf(101, 10) 10 >>> divf(101.1, 10) 10 >>> divf(4.5, 1) 4 """ return int(a // b) def floor(x, m=1): """ return largest multiple of , not greater than . """ # [[ m is 1 ]] gives a SyntaxWarning if m == 1: return intf(x) return m * int(x // m) def divc(a, b): """ ceiling division. always returns an int. >>> divc(100, 10) 10 >>> divc(101, 10) 11 >>> divc(101.1, 10) 11 >>> divc(4.5, 1) 5 """ return -int(-a // b) cdiv = divc def divr(a, b): """ round the value of a/b to the nearest integer. divr(a, b) = intr(fdiv(a, b)) >>> divr(0, 1) 0 >>> divr(5, 2) 3 >>> divr(10, -4) -3 """ if b < 0: (a, b) = (-a, -b) if a < 0: a = -a return -int((a + a + b) // (2 * b)) else: return int((a + a + b) // (2 * b)) def ceil(x, m=1): """ return lowest multiple of , not less than """ if m == 1: return intc(x) return m * -int(-x // m) def div(a, b): """ returns (a // b) if b exactly divides a, otherwise None >>> div(1001, 13) 77 >>> bool(div(101, 13)) False >>> div(42, 0) is None True """ if b == 0: return None (d, r) = divmod(a, b) if r != 0: return None return d def ediv(a, b): "return (a // b) if b exactly divides a, otherwise raise a ValueError." (d, r) = divmod(a, b) if r != 0: raise ValueError("inexact division") return d def fdiv(a, b, fn=float): """ float result of divided by . >>> fdiv(3, 2) 1.5 >>> fdiv(9, 3) 3.0 """ return fn(a) / fn(b) def is_duplicate(*s): """ check to see if arguments (as strings) contain duplicate characters. >>> is_duplicate("hello") True >>> is_duplicate("world") False >>> is_duplicate(99**2) False """ s = join(s) return len(set(s)) != len(s) # or using regexps #return True if re.search(r'(.).*\1', str(s)) else False duplicate = is_duplicate # iterative version of: # is_palindrome = lambda s: len(s) < 2 or (s[0] == s[-1] and is_palindrome(s[1:-1])) def is_palindrome(s): """ check to see if sequence is palindromic. >>> is_palindrome([1, 2, 3, 2, 1]) True >>> is_palindrome("ABBA") True >>> first(n for n in irange(0, inf) if not is_palindrome(nsplit(11**n))) [5] """ j = len(s) if j < 2: return True i = 0 j -= 1 while i < j: if s[i] != s[j]: return False i += 1 j -= 1 return True # is a number palindromic in base b def is_npalindrome(n, base=10): """ check if integer is palindromic in base . >>> is_npalindrome(123321) True >>> is_npalindrome(65535, base=16) True """ n = abs(n) if n % base == 0: return (n == 0) (a, b) = (n, 0) while a > b: (d, r) = divmod(a, base) b *= base b += r if a == b: return True a = d return a == b # originally called product(), but renamed to avoid name confusion with itertools.product() def multiply(seq, r=1, mod=None): """ return the product of the numeric sequence . if is specified this is used as the initial value of the product (and is the value returned when is empty). if is specified, the result at each stage is calculate mod . See also: math.prod() (Python 3.8+). >>> multiply(irange(1, 7)) 5040 >>> multiply([2] * 8) 256 >>> multiply(irange(100, 200), mod=1234) 18 """ if mod is None: for x in seq: r *= x else: for x in seq: r *= x r %= mod return r # vector dot product: dot(xs, ys, strict=0, fnp=multiply, fns=sum) def dot(*vs, **kw): """ this function takes a sequence of vectors provided as arguments, and calculates the product of the elements in the same position in each vector, and then sums these products. for two vectors this is the same as the vector dot product: dot((a1, a2, a3, ...), (b1, b2, b3, ...)) = a1 * b1 + a2 * b2 + a3 * b3 + ...) if the 'strict' argument is present it will be passed to zip() (which in supported Python versions will throw an error if the inputs are not of equal length), otherwise the length of vectors processed will be defined by the shortest input vector. the functions used for the product and sum functions can be defined with the parameters 'fnp' (default is: multiply) and 'fns' (default is: sum). >>> dot((1, 3, -5), (4, -2, -1)) 3 >>> call(dot, [(1, 3, -5)] * 2) 35 >>> call(dot, [(1, 3, -5)] * 3) -97 """ z = (zip(*vs, strict=kw['strict']) if 'strict' in kw else zip(*vs)) fns = kw.get('fns', sum) fnp = kw.get('fnp', multiply) return fns(map(fnp, z)) # multiple argument versions of basic operations: # - add, mul, bit_or, bit_and, bit_xor add = lambda *vs: sum(vs) mul = lambda *vs: multiply(vs) def bit_or(*vs): r = 0 for v in vs: r |= v return r def bit_xor(*vs): r = 0 for v in vs: r ^= v return r def bit_and(r, *vs): for v in vs: r &= v return r def _gcd(a, b): """ greatest common divisor (on positive integers). >>> gcd(123, 456) 3 >>> gcd(5, 7) 1 """ while b: (a, b) = (b, a % b) return a # or use math.gcd() [available from 3.5; from 3.9 = mgcd] gcd = getattr(math, 'gcd', _gcd) def _lcm(a, b): """ lowest common multiple (on positive integers). >>> lcm(123, 456) 18696 >>> lcm(5, 7) 35 """ return (a // gcd(a, b)) * b # or use math.lcm() [available from 3.5; from 3.9 = mlcm] lcm = getattr(math, 'lcm', _lcm) # Extended Euclidean Algorithm def egcd(a, b): """ Extended Euclidean Algorithm (on positive integers). returns integers (x, y, g) = egcd(a, b) where ax + by = g = gcd(a, b) Note that x and y are not necessarily positive integers. >>> egcd(120, 23) (-9, 47, 1) """ ## recursively... #if b == 0: return (1, 0, a) #(q, r) = divmod(a, b) #(s, t, g) = egcd(b, r) #return (t, s - q * t, g) # # or iteratively... (x0, x1) = (1, 0) (y0, y1) = (0, 1) while b: (q, r) = divmod(a, b) (a, b, x0, x1, y0, y1) = (b, r, x1, x0 - q * x1, y1, y0 - q * y1) return (x0, y0, a) # multiplicative inverse of mod def _invmod(n, m): """ return the multiplicative inverse of n mod m (or None if there is no inverse) i.e. the value x such that (n * x) % m = 1 e.g. the inverse of 2 (mod 9) is 5, as (2 * 5) % 9 = 1 >>> invmod(2, 9) 5 """ (x, y, g) = egcd(n, m) return ((x % m) if g == 1 else None) # from Python 3.8, pow() can do this for us if _pythonv > (3, 7): def invmod(n, m): return catch(pow, n, -1, m) invmod.__doc__ = _invmod.__doc__ else: invmod = _invmod # find square roots of mod # this is OK for relatively small m, but more efficient (and complex) # approaches are available (e.g. sympy.ntheory.sqrt_mod_iter) def sqrtmod(a, m): """ find square roots of a mod m. i.e. values x such that (x * x) is congurent to a (mod m). >>> sorted(sqrtmod(1, 16)) [1, 7, 9, 15] >>> sorted(sqrtmod(17, 43)) [19, 24] >>> sorted(sqrtmod(-1, 25)) [7, 18] """ a %= m for x in irange(0, m // 2): if (x * x) % m == a: # x is a root yield x # -x (mod m) is also a root if x > 0 and m > 2 * x: yield m - x # multiple GCD def mgcd(a, *rest): """ GCD of multiple (two or more) integers. See also: math.gcd() (Python 3.9+) >>> mgcd(123, 456) 3 >>> mgcd(123, 234, 345, 456, 567, 678, 789) 3 >>> mgcd(11, 37, 228) 1 >>> mgcd(56, 65, 671) 1 """ return reduce(gcd, rest, a) # multiple LCM def mlcm(a, *rest): """ LCM of multiple (two or more) integers. See also: math.lcm() (Python 3.9+) >>> mlcm(2, 3, 5, 9) 90 """ return reduce(lcm, rest, a) def is_coprime(*vs): return mgcd(*vs) == 1 # multiple divmod # hours, minutes, seconds: (h, m, s) = mdivmod(x, 60, 60) # days, hours, minutes, seconds: (d, h, m, s) = mdivmod(x, 24, 60, 60) # days, hours, minutes, seconds, fractional seconds: (d, h, m, s, f) = mdivmod(x, 24, 60, 60, 1) def mdivmod(x, *vs): rs = list() for v in reversed(vs): (x, r) = divmod(x, v) rs.insert(0, r) rs.insert(0, x) return rs # for those times when Rational() is overkill def fraction(a, b, *rest): """ return the numerator and denominator of the fraction a/b in lowest terms if more than 2 arguments are specified the sum of the arguments as (numerator, denominator) pairs is determined, so: fraction(a, b, c, d, e, f, ...) -> a/b + c/d + e/f + ... >>> fraction(286, 1001) (2, 7) >>> fraction(1, 2, 1, 3, 1, 6) # 1/2 + 1/3 + 1/6 = 1 (1, 1) >>> fraction(1, 2, 3, 4, 5, 6) # 1/2 + 3/4 + 5/6 = 25/12 (25, 12) """ if rest: for (c, d) in chunk(rest, 2): (a, b) = (a * d + b * c, b * d) if b < 0: (a, b) = (-a, -b) g = gcd(a, b) return (a // g, b // g) def format_fraction(n, d, base=10): s = int2base(n, base=base) if d == 1: return s return s + "/" + int2base(d, base=base) def ratio(*ns): """ return ratio of integers in in lowest terms. >>> ratio(6, 8) (3, 4) >>> ratio(6, 8, 10) (3, 4, 5) """ g = mgcd(*ns) return (ns if g == 1 else tuple(v // g for v in ns)) # import a value from a qualified spec, e.g.: # Q = import_fn('fractions.Fraction') # Q = import_fn('gmpy2.mpq') # Q = import_fn('mpmath.rational.mpq') # urlopen = import_fn('urllib2.urlopen') # Python 2 # urlopen = import_fn('urllib.request.urlopen') # Python 3 def import_fn(spec): # we could use importlib.import_module() here importer = lambda x: __import__(x, fromlist=['']) if '.' not in spec: return importer(spec) (mod, fn) = spec.rsplit('.', 1) return getattr(importer(mod), fn) # find an appropriate rational class # (could also try "sympy.Rational", but not for speed) # be aware when using gmpy2.mpq: # >> x = mpq(64) # >> y = mpq(x, 2) # will end up setting both x and y to 32 (they are the same object) # I have submitted a bug against gmpy2 (#334) # instead do this: # >> x = mpq(64) # >> y = x / 2 # if fix_gmpy2=1 is set, a workaround will be used for gmpy2 < 2.1.4 @static(src="gmpy2.mpq gmpy.mpq fractions.Fraction", impl=dict(), fix_gmpy2=None) def Rational(src=None, verbose=None): """ select a class for representing rational numbers. >> Q = Rational(verbose=1) [Rational: using gmpy2.mpq] >> Q = Rational(src="fractions.Fraction", verbose=1) [Rational: using fractions.Fraction] """ s = f = None if src is None: try: (s, f) = Rational.impl['*'] except KeyError: src = Rational.src if f is None: for s in src.split(): try: f = Rational.impl[s] break except KeyError: pass try: f = import_fn(s) except (ImportError, KeyError): continue Rational.impl[s] = f if '*' not in Rational.impl and src == Rational.src: Rational.impl['*'] = (s, f) # gmpy2 is fixed in v2.1.4+ if s == 'gmpy2.mpq' and Rational.fix_gmpy2 is None: Rational.fix_gmpy2 = (not catch(require, "gmpy2.version", "2.1.4")) break if verbose is None: verbose = ('v' in _PY_ENIGMA) if verbose: printf("[Rational: using {s}]", s=(s if f else f)) # fix for gmpy2.mpq() behaviour (issue #334) - may be fixed in gmpy2.version() > 2.1.2 if Rational.fix_gmpy2 and s == 'gmpy2.mpq': if verbose: printf("[Rational: applying fix for {s}]") f = (lambda x, y=None, fn=f: (fn(x) if y is None else fn(x) / y)) return f # create a function that will calculate a/b, and return an int if the result is an integer # or a rational object if the result is a rational class Rdiv(object): def __init__(self, F=None, src=None, verbose=None): self.F = F self.src = src self.verbose = verbose def __call__(self, a, b): (d, r) = divmod(a, b) if r == 0: return d if self.F is None: self.F = Rational(src=self.src, verbose=self.verbose) return self.F(a) / b # because mpq(x, y) changes x rdiv = Rdiv() @static(impl=None) def rational(*args, **kw): """ create an object representing a rational number. the class used is selected using Rational(), so for more control use: >> rational = Rational(verbose=1) or: >> rational = Rational(src="", verbose=1) to see what implementation is being used. (or set the 'v' flag in the environment variable PY_ENIGMA). >>> rational(1, 49) * 49 == 1 True """ if rational.impl is None: rational.impl = Rational() return rational.impl(*args, **kw) def factorial(a, *bs): """ return a! / b!. >>> factorial(6) 720 >>> factorial(10, 7) 720 number of anagrams of "mississippi" (len = 11; 4x i, 4x s, 2x p) >>> factorial(11, 4, 4, 2) 34650 """ if not bs: return math.factorial(a) r = None bs = sorted(bs, reverse=1) b = bs[0] if a - b < 100: r = multiply(irange(b + 1, a)) bs.pop(0) if r is None: r = math.factorial(a) for b in bs: if b == 1: break (r, z) = divmod(r, math.factorial(b)) if z != 0: raise ValueError("inexact division") return r # multinomial coefficient def multinomial(ks, n=None): """ calculate multinomial coefficient. e.g. number of anagrams of "mississippi" (len = 11; 4x i, 4x s, 2x p) >>> multinomial([4, 4, 2, 1]) 34650 >>> multinomial([4, 4, 2], 11) 34650 """ if n is None: n = sum(ks) return factorial(n, *ks) def nPr(n, r): """ permutations functions: n P r. the number of ordered r-length selections from n elements (elements can only be used once). See also: math.perm() (Python 3.8+). >>> nPr(10, 3) 720 """ if r > n: return 0 else: return math.factorial(n) // math.factorial(n - r) P = nPr def nCr(n, r): """ combinatorial function: n C r. the number of unordered r-length selections from n elements (elements can only be used once). See also: math.comb() (Python 3.8+). >>> nCr(10, 3) 120 """ if r > n: return 0 else: return math.factorial(n) // math.factorial(r) // math.factorial(n - r) C = nCr # NOTE: this corresponds to [[ select='R' ]] in subsets(), not [[ select='M' ]] def M(n, k): """ multichoose function: n M k. the number of unordered k-length selections from n elements where elements may be chosen multiple times. M(n, k) = icount(subsets(irange(1, n), size=k, select='R')) >>> M(10, 3) 220 """ return C(n + k - 1, k) Recurring = namedtuple('Recurring', 'i nr rr') def recurring(a, b, recur=0, base=10, digits=None): """ find recurring representation of the fraction / in the specified base. return strings (, , ) if you want rationals that normally terminate represented as non-terminating set >>> tuple(recurring(1, 7)) ('0', '', '142857') >>> tuple(recurring(3, 2)) ('1', '5', '') >>> tuple(recurring(3, 2, recur=1)) ('1', '4', '9') >>> tuple(recurring(5, 17, base=16)) ('0', '', '4B') """ # check input fraction if b == 0 or (recur and a == 0): raise ValueError("invalid input fraction: {a} / {b} [recur={recur}]".format(a=a, b=b, recur=bool(recur))) # sort out arguments neg = 0 if b < 0: (a, b) = (-a, -b) if a < 0: (a, neg) = (-a, 1) # the integer part (i, a) = divmod(a, b) if recur and a == 0: (i, a) = (i - 1, b) # record dividends r = dict() s = '' n = 0 while True: j = r.get(a) if j is not None: # have we had this dividend before? (i, nr, rr) = (int2base(i, base, digits=digits), s[:j], s[j:]) if neg and (nr or rr or i != '0'): i = '-' + i return Recurring(i, nr, rr) else: # no, we haven't r[a] = n n += 1 (d, a) = divmod(base * a, b) if recur and a == 0: (d, a) = (d - 1, b) if not (d == a == 0): # add to the digit string s += int2base(d, base, digits=digits) # Python 3.6: ...(*args, dp='.') def format_recurring(*args, **kw): """ format the (i, nr, rr) return from recurring() as a string >>> format_recurring(recurring(1, 7)) '0.(142857)...' >>> format_recurring(recurring(3, 2)) '1.5' >>> format_recurring(recurring(3, 2, recur=1)) '1.4(9)...' >>> format_recurring(recurring(5, 17, base=16)) '0.(4B)...' """ dp = kw.get('dp', '.') enc = kw.get('enc', '()') dots = kw.get('dots', '...') if len(args) == 1: args = args[0] (i, nr, rr) = args rr = (enc[0] + rr + enc[1] + dots if rr else '') return (i + dp + nr + rr if nr or rr else i) # recurring -> fraction def recurring2fraction(i, nr, rr, base=10, digits=None): """ turn the decimal representation .()... into a fraction in its lowest terms. >>> recurring2fraction('0', '', '142857') (1, 7) >>> recurring2fraction('1', '5', '') (3, 2) >>> recurring2fraction('1', '4', '9') (3, 2) >>> recurring2fraction('0', '', '4B', base=16) (5, 17) """ (p, q) = (len(nr), len(rr)) i = base2int(i, base=base, digits=digits) if q: if p: d = (base**(p + q)) - (base**p) n = base2int(nr, base=base, digits=digits) * ((base**q) - 1) + base2int(rr, base=base, digits=digits) else: d = (base**q) - 1 n = base2int(rr, base=base, digits=digits) elif p: d = base**p n = base2int(nr, base=base, digits=digits) else: return (i, 1) (a, b) = fraction(n, d) return ((i * b - a, b) if i < 0 else (i * b + a, b)) # see: Enigma 348 def reciprocals(k, b=1, a=1, m=1, M=inf, g=0, rs=[]): """ generate k whole numbers (d1, d2, ..., dk) such that 1/d1 + 1/d2 + ... + 1/dk = a/b the numbers are generated as an ordered list m = minimum allowed number M = maximum allowed number g = minimum allowed gap between numbers e.g. 3 reciprocals that sum to 1: 1/2 + 1/3 + 1/6 = 1 1/2 + 1/4 + 1/4 = 1 1/3 + 1/3 + 1/3 = 1 >>> list(reciprocals(3, 1)) [[2, 3, 6], [2, 4, 4], [3, 3, 3]] """ # are we done? if k == 1: (d, r) = divmod(b, a) if r == 0 and not (d < m or d > M): yield rs + [d] elif k == 2: # special case k = 2 if M == inf: dmin = divc(b + 1, a) else: dmin = divc(b * M, a * M - b) dmax = divf(b + b, a) for d in irange(max(m, dmin), min(M, dmax)): (e, r) = divmod(d * b, d * a - b) if r == 0 and not (e < d + g or e > M): yield rs + [d, e] else: if M == inf: # general case dmin = divc(b + 1, a) else: # but if M is given we can find a better dmin [suggested by frits] xs = list(M - g * i for i in xrange(k - 1)) xd = multiply(xs) xn = sum(xd // x for x in xs) dmin = divc(b * xd, a * xd - b * xn) dmax = divf(k * b, a) # find a suitable reciprocal for d in irange(max(m, dmin), min(M, dmax)): # solve for the remaining fraction [[Python 3: yield from ... ]] for ds in reciprocals(k - 1, b * d, a * d - b, d + g, M, g, rs + [d]): yield ds # command line arguments # fetch command line arguments from sys @static(argv=None) def get_argv(force=0, args=None): if force or get_argv.argv is None: get_argv.argv = (args if args is not None else sys.argv[1:]) return get_argv.argv # alias argv = get_argv # might have been better to use: arg(n, fn=identity, default=None, argv=None) # if 'p' is in PY_ENIGMA, then we will prompt # if 'v' is in PY_ENIGMA, then we will print values def arg(v, n, fn=identity, prompt=None, argv=None): """ if command line argument is specified return fn(argv[n]) otherwise return default value if argv is None (the default), then the value of sysv.argv[1:] is used >>> arg(42, 0, int, argv=['56']) 56 >>> arg(42, 1, int, argv=['56']) 42 """ if argv is None: argv = get_argv() r = (fn(argv[n]) if n < len(argv) else v) if 'p' in _PY_ENIGMA: if not prompt: prompt = "value" s = raw_input(sprintf("arg{n}: {prompt} [{r}] > ")).strip() if s: r = fn(s) if 'v' in _PY_ENIGMA: if not prompt: prompt = "value" printf("[arg{n}: {prompt} = {r!r}]") return r # get a list of similar arguments # if no arguments are collected the value of is returned def args(vs, n, fn=identity, prompt=None, argv=None): if argv is None: argv = get_argv() rs = (list(fn(v) for v in argv[n:]) or vs) if 'p' in _PY_ENIGMA: if not prompt: prompt = "values" s = raw_input(sprintf("args: {prompt} [{rs}] > ")).strip() if s: rs = list(fn(v) for v in re.split(r',\s*|\s+', s)) if 'v' in _PY_ENIGMA: if not prompt: prompt = "values" printf("[args: {prompt} = {rs!r}]") return rs # printf / sprintf variable interpolation # (see also the "say" module) # this works in all version of Python def __sprintf(fmt, vs): return fmt.format(**vs) # Python 3 has str.format_map(vs) def __sprintf3(fmt, vs): return fmt.format_map(vs) # in Python v3.6.x we are getting f"..." strings which can do this job # # NOTE: you lose the ability to do this: # # printf("... {d[x]} ...", d={ 'x': 42 }) -> "... 42 ..." # # instead you have to do this: # # printf("... {d['x']} ...", d={ 'x': 42 }) -> "... 42 ..." # # but you gain the ability to use arbitrary expressions: # # printf("... {a} + {b} = {a + b} ...", a=2, b=3) -> "... 2 + 3 = 5 ..." def __sprintf36(fmt, vs): return eval('f' + repr(fmt), vs) @static(fn=None) def _sprintf(fmt, vs, frame): # first try using the locals of the frame d = frame.f_locals if vs: d = update(d, vs) try: return _sprintf.fn(fmt, d) except (NameError, KeyError): pass # if that fails, try adding in the globals too d = update(frame.f_globals, frame.f_locals) if vs: d = update(d, vs) return _sprintf.fn(fmt, d) _sprintf.fn = __sprintf if _python > 2: _sprintf.fn = __sprintf3 if _pythonv > (3, 5): _sprintf.fn = __sprintf36 # print with variables interpolated into the format string def sprintf(fmt='', **kw): """ interpolate local variables and any keyword arguments into the format string . >>> (a, b, c) = (1, 2, 3) >>> sprintf("a={a} b={b} c={c}") 'a=1 b=2 c=3' >>> sprintf("a={a} b={b} c={c}", c=42) 'a=1 b=2 c=42' """ return _sprintf(fmt, kw, sys._getframe(1)) # print with local variables interpolated into the format string def printf(fmt='', **kw): """ print format string with interpolated local variables and keyword arguments. the final newline can be suppressed by ending the string with '\\'. >>> (a, b, c) = (1, 2, 3) >>> printf("a={a} b={b} c={c}") a=1 b=2 c=3 >>> printf("a={a} b={b} c={c}", c=42) a=1 b=2 c=42 """ s = _sprintf(fmt, kw, sys._getframe(1)) d = dict() # flush=1 if s.endswith('\\'): (s, d['end']) = (s[:-1], '') print(s, **d) def catch(fn, *args, **kw): """ evaluate the function with the given arguments, but if it throws an exception return None instead. >>> catch(divmod, 7, 0) is None True """ try: return fn(*args, **kw) except Exception: #print("catch: caught exception!") return # inclusive range iterator # TODO: should irange(4) = (1, 2, 3, 4) or (0, 1, 2, 3) ? # it's currently the latter, but maybe should be the former @static(inf=inf) # so b=irange.inf can be used def irange(a, b=None, step=1): """ irange(a, b) = a range iterator that includes both integer endpoints, and . if is specified as inf (or -inf for negative steps) the iterator will generate values indefinitely. irange(n) = if only one value is specified for the endpoints, then endpoints of 0 and (n - 1) are used (these are swapped if is negative), so that irange(n) produces n integers from 0 to n - 1. Note: Python's standard range iterator is available as xrange() if you want to emphasise the exclusion of the final endpoint. >>> list(irange(1, 9)) [1, 2, 3, 4, 5, 6, 7, 8, 9] >>> list(irange(9, 1, step=-1)) [9, 8, 7, 6, 5, 4, 3, 2, 1] >>> list(irange(0, 10, step=3)) [0, 3, 6, 9] >>> list(irange(10)) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] >>> list(irange(10, step=-1)) [9, 8, 7, 6, 5, 4, 3, 2, 1, 0] """ if step == 0: raise ValueError("irange: step cannot be 0") if b == inf: if step < 0: return xrange(0) elif b == -inf: if step > 0: return xrange(0) else: if b is None: (a, b) = ((0, a - 1) if step > 0 else (a - 1, 0)) return xrange(a, b + (1 if step > 0 else -1), step) return itertools.count(start=a, step=step) # inclusive range iterator that allows a fractional step def irangef(a, b, step=1): """ inclusive range iterator that allows the endpoints and the step to be fractional values. note that if float approximations are used for the step and/or endpoint then the final value may not be generated. >>> list(irangef(1, 2.5, step=0.5)) [1.0, 1.5, 2.0, 2.5] """ n = (inf if b == inf else divf(b - a, step)) for i in irange(0, n): yield a + i * step # flatten a list of lists def flatten(s, fn=list): """ flatten a list of lists (actually an iterator of iterators). the function: chain(*s) = flatten(s) is provided as a convenience. >>> flatten([[1, 2], [3, 4, 5], [6, 7, 8, 9]]) [1, 2, 3, 4, 5, 6, 7, 8, 9] >>> flatten(((1, 2), (3, 4, 5), (6, 7, 8, 9)), fn=tuple) (1, 2, 3, 4, 5, 6, 7, 8, 9) >>> flatten([['abc'], ['def', 'ghi']]) ['abc', 'def', 'ghi'] """ return fn(j for i in s if i is not None for j in i) # chain(a, b, c) = flatten([a, b, c]) # so: unpack(chain) = flatten def chain(*ss, **kw): """ a convenience function for calling flatten(): chain(a, b, c, ...) = flatten([a, b, c, ...], fn=iter) >>> chain("abc", (1, 2, 3), None, [4, 5, 6], fn=tuple) ('a', 'b', 'c', 1, 2, 3, 4, 5, 6) """ fn = kw.get("fn", iter) return flatten(ss, fn=fn) # interleave values from a bunch of iterators # flatten(zip(*ss), fn=iter) works if arguments are the same length def interleave(*ss, **kw): ss = list(iter(s) for s in ss) n = len(ss) while n > 0: i = 0 while i < n: try: yield next(ss[i]) i += 1 except StopIteration: ss.pop(i) n -= 1 # do we flatten this? def _flatten_test(s): # don't flatten strings if isinstance(s, (basestring, bytes)): return None # do flatten other sequences if isinstance(s, (Sequence, Iterable)): return s # otherwise don't flatten return None # a generator for flattening a sequence def _flattened(s, depth, test): d = (None if depth is None else depth - 1) for i in s: j = (test(i) if depth is None or depth > 0 else None) if j is None: yield i else: for k in _flattened(j, d, test): yield k # fully flatten a nested structure # ( has been renamed to ) def flattened(s, depth=None, test=_flatten_test, fn=None): """ fully flatten a nested structure (to depth , default is to fully flatten). can be used to determine how objects are flattened, it should return either - None, if the object is not to be flattened, or - an iterable of objects representing one level of flattening default behaviour is to flatten sequences other than strings >>> list(flattened([[1, [2, [3, 4, [5], [[]], [[6, 7], 8], [[9]]]], []]])) [1, 2, 3, 4, 5, 6, 7, 8, 9] >>> list(flattened([[1, [2, [3, 4, [5], [[]], [[6, 7], 8], [[9]]]], []]], depth=3)) [1, 2, 3, 4, [5], [[]], [[6, 7], 8], [[9]]] >>> list(flattened([['abc'], ['def', 'ghi']])) ['abc', 'def', 'ghi'] >>> flattened(42) 42 """ if test is None: test = fn z = (test(s) if depth is None or depth > 0 else None) if z is None: return s else: return _flattened(z, depth, test) # in Python 3 we could use functools.singledispatch # or in Python 3.10+ we could use structural pattern matching # to implement polymorphic functions (update(), delete(), append()) # however there does not seem to be a performance advantage # and it will fail with older versions of Python # return a copy of object , but with value at index for (k, v) in # can be a sequence of (k, v) pairs, or a sequence of keys, in which case # the values should be given in def update(s, ps=(), vs=None): """ create an updated version of object which is the same as except that the value at index is for the keys and values provided in and . can either be a sequence of (, ) pairs, or can be a sequence of keys and the corresponding sequence of values. >>> update([0, 1, 2, 3], [(2, 'foo')]) [0, 1, 'foo', 3] >>> update(dict(a=1, b=2, c=3), 'bc', (4, 9)) == dict(a=1, b=4, c=9) True >>> update((1, 2, 3), [(2, 4)]) (1, 2, 4) """ if vs is not None: ps = zip(ps, vs) # allow updating of immutable types: tuple, string fn = None if isinstance(s, tuple): fn = type(s) s = list(s) elif isinstance(s, basestring): fn = ''.join s = list(s) else: # use copy() method if available # otherwise create a new object initialised from the old one cpy = getattr(s, 'copy', None) s = (cpy() if cpy else type(s)(s)) # use update() method if available upd = getattr(s, 'update', None) if upd: upd(ps) else: # otherwise update the pairs individually for (k, v) in ps: s[k] = v # return the new object return (fn(s) if fn else s) # return a copy of object with values at indices removed def delete(s, ks=()): """ return an updated version of object with items at keys removed. >>> delete(dict(a=1, b=2, c=3), 'bc') == dict(a=1) True >>> delete("bananas", [0, 2, 4, 6]) 'aaa' """ fn = None if isinstance(s, list): s = list(s) ks = sorted(ks, reverse=1) elif isinstance(s, tuple): fn = type(s) s = list(s) ks = sorted(ks, reverse=1) elif isinstance(s, basestring): fn = ''.join s = list(s) ks = sorted(ks, reverse=1) else: # use copy() method if available # otherwise create a new object initialised from the old one cpy = getattr(s, 'copy', None) s = (cpy() if cpy else type(s)(s)) # remove specified keys for k in ks: del s[k] # return the new object return (fn(s) if fn else s) # this unifies adding an element (or elements) to a container: # - string: '123' + '4' -> append('123', '4') # - list: [1, 2, 3] + [4] -> append([1, 2, 3], 4) # - tuple: (1, 2, 3) + (4,) -> append((1, 2, 3), 4) # - set: {1, 2, 3} | {4} -> append({1, 2, 3}, 4) def append(s, *vs): """ make a new container, the same as but with additional values added. >>> append((1, 2, 3), 4) (1, 2, 3, 4) >>> append([1, 2, 3], 4) [1, 2, 3, 4] >>> append({1, 2, 3}, 4) == {1, 2, 3, 4} True >>> append('123', '4') '1234' """ if isinstance(s, list): s = type(s)(s) s.extend(vs) return s if isinstance(s, tuple): return s + vs if isinstance(s, basestring): return s + str.join('', vs) if isinstance(s, set): s = s.copy() s.update(vs) return s if isinstance(s, frozenset): return s.union(vs) if isinstance(s, multiset): return s.copy().update_from_seq(vs) raise ValueError(sprintf("append() can't handle container of type {x}", x=type(s))) # adjacency matrix for an n (columns) x m (rows) grid # entries are returned as lists in case you want to modify them before use def grid_adjacency(n, m, deltas=None, include_adjacent=1, include_diagonal=0, include_self=0, fn=None): """ this function generates the adjacency matrix for a grid with n columns and m rows, represented by a linear array of size n * m. the element in the (i, j)th position in the grid is at index (i + n * j) in the array. it returns an array, where the entry at index k is the collection of indices into the linear array that are adjacent to the square at index k. if 'fn' is specified then it is used to collect the indices, otherwise they are returned as a list. the default behaviour is to treat the squares immediately N, S, E, W of the target square as being adjacent, although this can be controlled with the 'deltas' parameter, it can be specified as a list of (x, y) deltas to use instead. if 'deltas' is not specified the 'include_adjacent', 'include_diagonal' and 'include_self' flags are used to specify which squares are adjacent to the target square: 'include_adjacent' includes the N, S, E, W squares 'include_diagonal' includes the NW, NE, SW, SE squares 'include_self' includes the square itself >>> grid_adjacency(2, 2) [[1, 2], [0, 3], [0, 3], [1, 2]] >>> sorted(grid_adjacency(3, 3)[4]) [1, 3, 5, 7] >>> sorted(grid_adjacency(3, 3, include_diagonal=1)[4]) [0, 1, 2, 3, 5, 6, 7, 8] """ # if deltas aren't provided use standard deltas if deltas is None: deltas = list() if include_adjacent: deltas.extend([(0, -1), (-1, 0), (1, 0), (0, 1)]) if include_diagonal: deltas.extend([(-1, -1), (1, -1), (-1, 1), (1, 1)]) if include_self: deltas.append((0, 0)) # construct the adjacency matrix t = n * m r = [None] * t for y in xrange(0, m): for x in xrange(0, n): s = list() for (dx, dy) in deltas: (x1, y1) = (x + dx, y + dy) if not (x1 < 0 or y1 < 0 or x1 + 1 > n or y1 + 1 > m): s.append(x1 + y1 * n) r[x + y * n] = (fn(s) if fn else s) return r # cumulative sum def csum(seq, s=0, fn=operator.add, empty=0): """ generate cumulative partial sums from sequence . 's' is the initial value, and 'fn' is the function used to update it with each element of the sequence. if 'empty' is set to a true value then the initial value 's' will be initially returned. See also: itertools.accumulate() (Python 3.2+). >>> list(csum(irange(1, 10))) [1, 3, 6, 10, 15, 21, 28, 36, 45, 55] >>> list(csum(irange(1, 10), fn=operator.mul, s=1)) [1, 2, 6, 24, 120, 720, 5040, 40320, 362880, 3628800] """ if empty: yield s for x in seq: s = fn(s, x) yield s # cumulative slice def cslice(seq, empty=0): """ generate an iterator that is the cumulative slices of a sequence. >>> list(cslice([1, 2, 3])) [[1], [1, 2], [1, 2, 3]] >>> list(cslice('python')) ['p', 'py', 'pyt', 'pyth', 'pytho', 'python'] """ for i in irange((0 if empty else 1), len(seq)): yield seq[:i] # overlapping tuples from a sequence def tuples(seq, n=2, circular=0, fn=tuple): """ generate overlapping -tuples from sequence . (for non-overlapping tuples see chunk()). if 'circular' is set to true, then values from the beginning of will be used to complete tuples when the end is reached. See also: itertools.pairwise() (Python 3.10+). >>> list(tuples('ABCDE')) [('A', 'B'), ('B', 'C'), ('C', 'D'), ('D', 'E')] >>> list(tuples(irange(1, 5), 3)) [(1, 2, 3), (2, 3, 4), (3, 4, 5)] >>> list(tuples(irange(1, 5), 3, circular=1)) [(1, 2, 3), (2, 3, 4), (3, 4, 5), (4, 5, 1), (5, 1, 2)] """ assert n > 0 i = iter(seq) if circular and n > 1: # we need extract the first (n - 1) items and add them to the end xs = first(i, n - 1) if not xs: return m = len(xs) if m < n - 1: i = itertools.chain(xs, i, (xs[k % m] for k in xrange(n - 1))) else: i = itertools.chain(xs, i, xs) t = list() try: # collect the first tuple for _ in xrange(n): t.append(next(i)) while True: # return the tuple yield fn(t) # move the next value in to the tuple t.append(next(i)) t.pop(0) except StopIteration: return def last(seq, count=1, fn=list): """ find the last items in sequence . >>> last([1, 2, 3, 4]) [4] >>> last(Primes(30), 3) [19, 23, 29] """ try: x = seq[-count:] if len(x) < count: return except TypeError: x = None for x in tuples(seq, count, fn=list): pass if x is None: return return (x if fn == list else fn(x)) def contains(seq, subseq): """ return the position in that occurs as a contiguous subsequence or -1 if it is not found >>> contains("abcdefghijkl", "def") 3 >>> contains("abcdefghijkl", "hik") -1 >>> contains(primes, [11, 13, 17, 19]) 4 >>> contains([1, 2, 3], [1, 2, 3]) 0 >>> contains([1, 2, 3], []) 0 """ subseq = tuple(subseq) n = len(subseq) if n == 0: return 0 k = 0 i = 1 for x in seq: if x == subseq[k]: k += 1 if k == n: return i - n else: k = 0 i += 1 return -1 # subseqs: generate the subsequences of an iterator -> replaced by subsets() # bit permutations # see: https://enigmaticcode.wordpress.com/2017/05/20/bit-twiddling/ def bit_permutations(a, b=None): """ generate numbers in order that have the same number of bits set in their binary representation. numbers start at and are generated while they are smaller than . to generate all numbers with k bits start start with: a = pow(2, k) - 1 a = (1 << k) - 1 >>> list(bit_permutations(3, 20)) [3, 5, 6, 9, 10, 12, 17, 18] >>> first(bit_permutations(1), 11) [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] """ if a == 0: yield a return while b is None or a < b: yield a t = (a | (a - 1)) + 1 a = t | ((((t & -t) // (a & -a)) >> 1) - 1) def bit_from_positions(ps): """ construct an integer with bit positions in set >>> bit_from_positions({1, 3, 7, 11}) 2186 """ return bit_or(*(1 << p for p in ps)) def bit_positions(x): """ return the positions of bits set in integer >>> list(bit_positions(2186)) [1, 3, 7, 11] """ i = 0 while x: if x & 1: yield i x >>= 1 i += 1 # for "coin puzzles", see also: Denominations() # simple express: def express(t, ds, qs=None, min_q=0): """ express total using denominations . optional: using quantities chosen from or: minimum quantity (non-negative integer) and should be increasing sequences. generated values are the quantities for each denomination in . >>> list(express(20, (3, 5, 7))) [[0, 4, 0], [1, 2, 1], [2, 0, 2], [5, 1, 0]] >>> list(express(20, (3, 5, 7), min_q=1)) [[1, 2, 1]] >>> list(express(20, (3, 5, 7), qs=(0, 1, 2))) [[1, 2, 1], [2, 0, 2]] the number of ways to change 1 pound into smaller coins >>> icount(express(100, (1, 2, 5, 10, 20, 50))) 4562 """ ds = list(ds) assert ds and ds[0] > 0, sprintf("invalid denominations {ds}") if qs: return express_quantities(t, ds, qs) if min_q > 0: return express_denominations_min(t, ds, min_q) else: return express_denominations(t, ds) # express total using denominations def express_denominations(t, ds, ss=[]): if t == 0: if not ds: yield ss else: yield ss + [0] * len(ds) elif ds: d = ds[0] (k, r) = divmod(t, d) if len(ds) == 1: if r: return qs = [k] else: qs = irange(0, k) for q in qs: for r in express_denominations(t - d * q, ds[1:], ss + [q]): yield r # express total using denominations , min quantity def express_denominations_min(t, ds, min_q): # allocate the minimum quantities t -= min_q * sum(ds) if t == 0: yield [min_q] * len(ds) elif t > 0: # solve for the remaining amount for ss in express_denominations(t, ds): # add in the initial quantities yield list(q + min_q for q in ss) # express total using denominations , quantities chosen from def express_quantities(t, ds, qs, ss=[]): if t == 0: if not ds: yield ss elif 0 in qs: yield ss + [0] * len(ds) elif ds: d = ds[0] for q in qs: if d * q > t: break for r in express_quantities(t - d * q, ds[1:], qs, ss + [q]): yield r # An implementation of the Boecker-Liptak Money Changing algorithm from: # # https://bio.informatik.uni-jena.de/bib2html/downloads/2007/BoeckerLiptak_FastSimpleAlgorithmMoneyChanging_Algorithmica_2007.pdf # https://pdfs.semanticscholar.org/14ac/14a15ebc31b58a4ac04328f9824f743a1e4e.pdf # # this implementation uses the "Round Robin" algorithm before optimisations. # make the Extended Residue Table def _residues(vs): # empty table res = [None] * len(vs) # initial row v0 = vs[0] r = [inf] * v0 r[0] = 0 # fill out each row for (i, vi) in enumerate(vs): if i > 0: d = gcd(v0, vi) for p in xrange(d): m = min(r[q] for q in xrange(p, v0, d)) if m < inf: for c in xrange(v0 // d): m += vi j = m % v0 if r[j] < m: m = r[j] else: r[j] = m res[i] = list(r) return res # generate possible expressions of t def _find_all(t, vs, i, c, res): v0 = vs[0] if i == 0: c[0] = t // v0 yield tuple(c) else: vi = vs[i] m = lcm(v0, vi) d = m // vi for j in xrange(0, d): c[i] = j u = t - j * vi b = res[i - 1][u % v0] while u >= b: for x in _find_all(u, vs, i - 1, c, res): yield x u -= m c[i] += d class Denominations(object): """ An implementation of the Boecker-Liptak Money Changing algorithm. The denominations passed in are sorted into increasing order, and accessible via the 'denominations' attribute. Quantities returned by 'express()' are in the same order as the 'denominations' attribute. >>> sorted(Denominations([3, 5, 7]).express(20)) [(0, 4, 0), (1, 2, 1), (2, 0, 2), (5, 1, 0)] >>> sorted(Denominations([3, 5, 7]).express(20, min_q=1)) [(1, 2, 1)] the number of ways to change 1 pound into smaller coins: >>> Denominations([1, 2, 5, 10, 20, 50]).count(100) 4562 using at least 1 of each type of coin: >>> Denominations([1, 2, 5, 10, 20, 50]).count(100, min_q=1) 15 the largest non-McNugget number: >>> Denominations([6, 9, 20]).frobenius() 43 """ def __init__(self, *denominations): # preferred initialisation is to pass a sequence of denominations if len(denominations) == 1: denominations = denominations[0] # first sort the denominations try: ds = tuple(sorted(denominations)) except TypeError: ds = () if not (len(ds) > 1 and ds[0] > 0 and seq_all_different(ds)): raise ValueError(sprintf("invalid denominations: {denominations}")) self.denominations = ds # compute the extended residue table for the given denominations self.residues = _residues(ds) # generate different ways to express def express(self, amount, min_q=0): """ generate the different ways to express the given amount. if min_q is specified (non-negative integer), at least that many instances of each denomination must be used. """ n = len(self.denominations) if min_q == 0: for t in _find_all(amount, self.denominations, n - 1, [0] * n, self.residues): yield t elif min_q > 0: amount -= min_q * sum(self.denominations) if amount == 0: yield (min_q,) * n elif amount > 0: for t in _find_all(amount, self.denominations, n - 1, [0] * n, self.residues): yield tuple(x + min_q for x in t) # count the number of ways to express def count(self, amount, min_q=0): """count the number of ways of expressing an amount""" return icount(self.express(amount, min_q=min_q)) # return the Frobenius number (the largest amount that cannot be changed) def frobenius(self): """return the largest amount not expressible using the denominations""" m = max(self.residues[-1]) return (None if m == inf else m - self.denominations[0]) # return a function to generate k-sequences of positive integers with a particular total def Decompose(k=None, increasing=1, sep=1, min_v=1, max_v=inf, fn=identity): """ return a function to generate k-sequences of non-negative integers that sum to a chosen total k = length of sequences to generate increasing = +1 (increasing sequences [default]); -1 (decreasing sequences); or 0 sep = separation between numbers; 0 allows repeats [default: 1] min_v = minimum permissible value (non-negative integer) max_v = maximum permissible value (non-negative integer, or inf) fn = return type (default is to return tuples) """ # decompose t into k increasing numbers, in range [min_v, max_v] # d = delta between numbers (for inc/dec seqs) # R = function to calculate remaining values # M = function to calculate next minimum value # r = reverse return values # fn = return type # ns = numbers collected so far def _decompose(t, k, min_v, max_v, d, R, M, r, fn, ns=()): if k == 0: if t == 0: yield fn(()) elif k == 1: if not (t < min_v or t > max_v): ns += (t,) yield fn(ns[::-1] if r else ns) else: k_ = k - 1 for n in irange(min_v, min(max_v, R(t, k, k_, min_v))): for z in _decompose(t - n, k_, M(n, d), max_v, d, R, M, r, fn, ns + (n,)): yield z if increasing == 0: # generate increasing sequences with the appropriate sep value # and then permute the answers (which may contain repeats if sep=0) f = Decompose(k, increasing=1, sep=sep, min_v=min_v, max_v=max_v, fn=fn) perm = (mpermutations if sep == 0 else itertools.permutations) return (lambda t, k=k: flatten((perm(ns, k) for ns in f(t, k)), fn=iter)) else: d = abs(sep) if d == 0: R = (lambda t, k, k_, m: t // k) M = (lambda n, d: n) elif d == 1: R = (lambda t, k, k_, m: (t - (k * k_) // 2) // k) M = (lambda n, d: n + 1) else: R = (lambda t, k, k_, m: (t - (d * k * k_) // 2) // k) M = (lambda n, d: n + d) r = (increasing < 0) return (lambda t, k=k, min_v=min_v: _decompose(t, k, min_v, max_v, d, R, M, r, fn)) # all-in-one def decompose(t, k, increasing=1, sep=1, min_v=1, max_v=inf, fn=identity): """ decompose in -sequences of non-negative integers that sum to t = total sum of each sequence k = length of sequences to generate increasing = +1 (increasing sequences); -1 (decreasing sequences); or 0 sep = separation between numbers (if increasing != 0); 0 allows repeats min_v = minimum permissible value (non-negative integer) max_v = maximum permissible value (non-negative integer, or inf) fn = return type (default is to return tuples) >>> sorted(decompose(10, 3, increasing=1, min_v=1)) [(1, 2, 7), (1, 3, 6), (1, 4, 5), (2, 3, 5)] >>> sorted(decompose(8, 3, increasing=1, min_v=0)) [(0, 1, 7), (0, 2, 6), (0, 3, 5), (1, 2, 5), (1, 3, 4)] >>> sorted(decompose(8, 3, increasing=1, sep=0, min_v=1)) [(1, 1, 6), (1, 2, 5), (1, 3, 4), (2, 2, 4), (2, 3, 3)] >>> sorted(decompose(5, 3, increasing=0, sep=0, min_v=1)) [(1, 1, 3), (1, 2, 2), (1, 3, 1), (2, 1, 2), (2, 2, 1), (3, 1, 1)] """ return call(Decompose(increasing=increasing, sep=sep, min_v=min_v, max_v=max_v, fn=fn), (t, k)) ############################################################################### # exact set cover (using Knuth's Algorithm X) # in-place algorithmX implementation (X, soln are modified) def algorithmX(X, Y, soln): if not X: yield soln else: c = min(X.keys(), key=(lambda k: len(X[k]))) # copy X[c], as X is modified (could use sorted(X[c]) for stability) for r in list(X[c]): soln.append(r) # cols = select(X, Y, r) cols = list() for j in Y[r]: for i in X[j]: for k in Y[i]: if k != j: X[k].remove(i) cols.append(X.pop(j)) for z in algorithmX(X, Y, soln): yield z # deselect(X, Y, r, cols) for j in reversed(Y[r]): X[j] = cols.pop() for i in X[j]: for k in Y[i]: if k != j: X[k].add(i) soln.pop() # input: ss = sequence of collections of sets [ [a0, a1, ...], [b1, b2, ...], [c1, c2, ...] ... ] # output: sequence of sets (a, b, c, ...) one from each collection def exact_cover(sss, tgt=None): """ given a collection of sets (of sets): [a0, a1, ...] [b1, b2, ...] ... an exact cover is a selection of sets: [a, b, ...] where a is chosen from the first collection, b from the second, etc. and each element of the target set appears in exactly one of the sets in the cover. if the target set is not specified, it is the collection of all elements contained in any of the provided sets. """ # map elements to indices if tgt is None: tgt = union(union(ss) for ss in sss) tgt = sorted(tgt) n = len(tgt) m = dict((x, i) for (i, x) in enumerate(tgt)) # set up Y, one row for each position Y = list() for (j, ss) in enumerate(sss, start=n): for s in ss: y = list(m[x] for x in s) y.append(j) Y.append(y) # set up X as a dict of sets X = dict((k, set()) for k in irange(0, j)) for (i, y) in enumerate(Y): for k in y: X[k].add(i) # find exact covers using algorithmX k = len(sss) for rs in algorithmX(X, Y, list()): # turn the selected rows of Y, back into sets r = [None] * k for i in rs: y = Y[i] r[y[-1] - n] = list(tgt[i] for i in y[:-1]) yield r # exact multiset cover (see: Enigma 1712, Teaser 2690) def mcover(m, tgt, reject=None): """ find exact multiset covers. is a map of keys to multisets of values is a multiset of values to target is a function that can be used to reject partial solutions solutions are returned as an (unordered) list of keys, where the combined multisets of values corresponding to those keys give exactly the target multiset. """ # a variation on Knuth's Algorithm X def _mcover(m, tgt, X, ss): # are we done? if not tgt: yield ss else: # choose a column to work on c = min(tgt.keys(), key=(lambda k: (len(X[k]), -tgt[k]))) # consider subsets with this value xs = sorted(X[c]) for (j, n) in enumerate(xs, start=1): s = m[n] # update the target (tgt_, rs) = tgt.differences(s) if rs: continue # is this sequence acceptable? ss_ = ss + [n] if reject and reject(ss_): continue # remove the value (and any prior values) from consideration discard = set(xs[:j]) # remove any values in columns that have reached 0 for (k, v) in X.items(): if v and tgt_.count(k) == 0: discard.update(v) # recurse with a new target and X X_ = dict((k, v.difference(discard)) for (k, v) in X.items()) for z in _mcover(m, tgt_, X_, ss_): yield z # X tells us what elements of the target are involved in which values if not isinstance(tgt, multiset): tgt = multiset(tgt) ks = set(tgt.keys()) X = dict((k, set()) for k in ks) for (v, s) in m.items(): if ks.issuperset(s.keys()): for k in s.keys(): X[k].add(v) # check each target element appears if not all(X.values()): return () # solve using the variation on Algorithm X return _mcover(m, tgt, X, list()) ############################################################################### # Numerical approximations # a simple record type class for results # (Python 3.3 introduced types.SimpleNamespace) class Record(object): # best called as Record.update(self, ...) def update(self, **vs): """update values in a record""" self.__dict__.update(vs) return self # __init__ is the same as update def __init__(self, **vs): Record.update(self, **vs) def __iter__(self): d = self.__dict__ for k in sorted(d.keys()): yield (k, d[k]) def __repr__(self): return self.__class__.__name__ + map2str((k, repr(v)) for (k, v) in self) # best called as Record.map(self, ...) def map(self): return self.__dict__ # a golden section search minimiser # f - function to minimise # a, b - bracket to search # t - tolerance # m - metric def gss_minimiser(f, a, b, t=1e-9, m=None): # apply any metric fn = (f if m is None else (lambda x: m(f(x)))) R = 0.5 * (math.sqrt(5.0) - 1.0) C = 1.0 - R (x1, x2) = (R * a + C * b, C * a + R * b) (f1, f2) = (fn(x1), fn(x2)) while b - a > t: if f1 > f2: (a, x1, f1) = (x1, x2, f2) x2 = C * a + R * b f2 = fn(x2) else: (b, x2, f2) = (x2, x1, f1) x1 = R * a + C * b f1 = fn(x1) v = (x1 if f1 < f2 else x2) return Record(v=v, fv=f(v), t=t) find_min = gss_minimiser find_min.__name__ = 'find_min' find_min.__doc__ = """ find the minimum value of a (well behaved) function over an interval. f = function to minimise (should take a single float argument) a, b = the interval to minimise over (a < b) t = the tolerance to work to m = the metric we want to minimise (default is None = the value of the function) the result is returned as a record with the following fields: v = the calculated value at which the function is minimised fv = the value of the function at v t = the tolerance used >>> r = find_min(lambda x: sq(x - 2), 0.0, 10.0) >>> round(r.v, 6) 2.0 """ # NOTE: using functools.partial and setting __name__ and __doc__ doesn't work (in Python 2.7 and 3.3) # see: http://bugs.python.org/issue12790 def find_max(f, a, b, t=1e-9): """ find the maximum value of a (well behaved) function over an interval. f = function to maximise (should take a single float argument) a, b = the interval to search (a < b) t = the tolerance to work to the result is returned as a record with the following fields: v = the calculated value at which the function is maximised fv = the value of the function at v t = the tolerance used >>> r = find_max(lambda x: 9 - sq(x - 2), 0.0, 10.0) >>> round(r.v, 6) 2.0 """ return gss_minimiser(f, a, b, t=t, m=(lambda x: -x)) def find_zero(f, a, b, t=1e-9, ft=1e-6): """ find the zero of a (well behaved) function over an interval. f = function to find the zero of (should take a single float argument) a, b = the interval to search (a < b) t = the tolerance to work to the result is returned as a record with the following fields: v = the calculated value at which the function is zero fv = the value of the function at v t = the tolerance used >>> r = find_zero(lambda x: sq(x) - 4, 0.0, 10.0) >>> round(r.v, 6) 2.0 >>> r = find_zero(lambda x: sq(x) + 4, 0.0, 10.0) # doctest: +IGNORE_EXCEPTION_DETAIL Traceback (most recent call last): ... ValueError: Value not found """ r = find_min(f, a, b, t, m=abs) if ft < abs(r.fv): raise ValueError("Value not found") r.ft = ft return r def find_value(f, v, a, b, t=1e-9, ft=1e-6): """ find the value of a (well behaved) function over an interval. f = function to find the value of (should take a single float argument) a, b = the interval to search (a < b) t = the tolerance to work to the result is returned as a record with the following fields: v = the calculated value at which the function is the specified value fv = the value of the function at v t = the tolerance used >>> r = find_value(lambda x: sq(x) + 4, 8.0, 0.0, 10.0) >>> round(r.v, 6) 2.0 """ r = find_zero((lambda x: f(x) - v), a, b, t, ft) r.fv += v return r def line_intersect(p1, p2, p3, p4, internal=0, div=fdiv): """ Find the intersection of 2 lines defined by points: line 1 passes through p1 and p2 (= (x1, y1) and (x2, y2)) line 2 passes through p3 and p4 (= (x3, y3) and (x4, y4)) internal can be set to: 1, 2, 1+2 to check the intersection is internal to the specified line segments, if not an exception is raised div is set to an appropriate division function (default is fdiv() for floats, but the result of Rational() could be used) return value is a Record object: pt = (x, y) value of intersection x = x y = y q1 = fraction along line 1 (0 = p1, 1 = p2) q2 = fraction along line 2 (0 = p3, 1 = p4) """ ((x1, y1), (x2, y2), (x3, y3), (x4, y4)) = (p1, p2, p3, p4) (dx21, dx43, dx31) = (x2 - x1, x4 - x3, x3 - x1) (dy21, dy43, dy31) = (y2 - y1, y4 - y3, y3 - y1) z = dx21 * dy43 - dy21 * dx43 if z == 0: raise ValueError("invalid lines") # calculate parameter for line 1 = (x1 + r1(x2 - x1), y1 + r1(y2 - y1)) q1 = div(dx31 * dy43 - dy31 * dx43, z) if internal & 1 and (q1 < 0 or q1 > 1): raise ValueError("external intersection on line 1") # calculate intersection point (x, y) = (x1 + q1 * dx21, y1 + q1 * dy21) # calculate parameter for line 2 = (x3 + q2(x4 - x3), y3 + q2(y4 - x3)) if abs(dx43) > abs(dy43): q2 = div(x - x3, dx43) else: q2 = div(y - y3, dy43) if internal & 2 and (q2 < 0 or q2 > 1): raise ValueError("external intersection on line 2") # return intersection point (pt) (also: x, y, q1, q2) return Record(pt=(x, y), x=x, y=y, q1=q1, q2=q2) # return a line segment the same length as (p1, p2) that is its perpendicular bisector # i.e. (p1, p2) forms the diagonal of a square, the other diagonal is the return value (p3, p4) def line_bisect(p1, p2, div=fdiv): """ Return a line segment that is a perpendicular bisector of the line segment defined by p1 and p2 (= (x1, y1) and (x2, y2)). The value returned (p3, p4) (= (x3, y3), (x4, y4)) is a line segment that forms a diagonal of a square, where the other diagonal is (p1, p2). """ ((x1, y1), (x2, y2)) = (p1, p2) s = fdiv(x1 + x2 + y1 + y2, 2) return ((s - y1, s - x2), (s - y2, s - x1)) ############################################################################### # Roman Numerals # the following is good for numbers in the range 1 - 4999 _romans = ( # X bar = 10000 # V bar = 5000 # (, , ) ('M', 1000, 4), ('CM', 900, 1), ('D', 500, 1), ('CD', 400, 1), ('C', 100, 3), ('XC', 90, 1), ('L', 50, 1), ('XL', 40, 1), ('X', 10, 3), ('IX', 9, 1), ('V', 5, 1), ('IV', 4, 1), ('I', 1, 4) ) def int2roman(x): """ return a representation of an integer (from 1 to 4999) as a Roman Numeral >>> int2roman(4) 'IV' >>> int2roman(1999) 'MCMXCIX' """ x = as_int(x) if not (0 < x < 5000): raise ValueError("integer out of range: {x}".format(x=x)) s = list() for (n, i, m) in _romans: (d, r) = divmod(x, i) if d < 1: continue s.append(n * d) x = r return join(s) def roman2int(x): """ return the integer value of a Roman Numeral . >>> roman2int('IV') 4 >>> roman2int('IIII') 4 >>> roman2int('MCMXCIX') 1999 """ x = str(x).upper() p = r = 0 for (n, i, m) in _romans: (l, c) = (len(n), 0) while x[p:p + l] == n and c < m: r += i p += l c += 1 if p < len(x): raise ValueError("invalid Roman numeral: {x}".format(x=x)) return r def is_roman(x): """ check if a Roman Numeral is valid. >>> is_roman('IV') True >>> is_roman('IIII') True >>> is_roman('XIVI') False """ x = str(x).upper() if x == 'IIII': return True try: i = roman2int(x) except ValueError: return False return int2roman(i) == x # (default) digits for use in converting bases _DIGITS = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ' @static(digits=_DIGITS) def base_digits(*args): """ get or set the default string of digits used to represent numbers. with no arguments the current string of digits is returned. with an argument the current string is set, and the new string returned, if the argument is None (or any non-True value) then the standard default string is used (0-9 then A-Z). NOTE: this is a global setting and will affect all subsequent base conversions. """ assert len(args) < 2 if args: base_digits.digits = (args[0] or _DIGITS) return base_digits.digits def int2base(i, base=10, width=None, pad=None, group=None, sep=",", digits=None): """ convert an integer to a string representation in the specified base . if the parameter is specified the number of digits will be padded to value of using the character. if is positive pad characters will be added on the left, if negative they are added on the right. The default pad character is the digit 0. if the parameter is specified the digits are grouped into blocks of digits and separated by the string (this happens after the digits are padded to any specified ). if is positive the groups start from the right, if negative they start from the left. By default this routine only handles single digits up 36 in any given base, but the parameter can be specified to give the symbols for larger bases. And if the error=1 parameter is given then invalid digits will be substituted with"", where n is the digit value in decimal. >>> int2base(-42) '-42' >>> int2base(3735928559, base=16) 'DEADBEEF' >>> int2base(-3735928559, base=16, digits='0123456789abcdef') '-deadbeef' >>> int2base(190, base=3, digits='zyx') 'xyzzy' >>> int2base(29234652, base=36) 'HELLO' >>> int(int2base(123456, base=14), base=14) 123456 >>> int2base(84, base=2, width=9, group=3, sep="_") '001_010_100' """ assert base > 1, "invalid base {base}".format(base=base) if digits is None: digits = base_digits() (p, r) = ('', None) if i < 0: (p, i) = ('-', -i) # if there aren't enough digits switch to {::...} format if len(digits) < base: r = join(nsplit(i, base=base), sep=':', enc="{}") elif i == 0: r = digits[0] else: r = list() while i > 0: (i, n) = divmod(i, base) r.insert(0, digits[n]) r = join(r) if width is not None: if pad is None: pad = digits[0] r = (r.rjust(width, pad) if width > 0 else r.ljust(-width, pad)) if group is not None: (s, group) = ((-1, group) if group > 0 else (1, -group)) r = join((join(x) for x in chunk(r[::s], group)), sep=sep[::s])[::s] return (p + r if p else r) def base2int(s, base=10, strip=0, digits=None): """ convert a string representation of an integer in the specified base to an integer. >>> base2int('-42') -42 >>> base2int('xyzzy', base=3, digits='zyx') 190 >>> base2int('HELLO', base=36) 29234652 """ assert base > 1, "invalid base {base}".format(base=base) if digits is None: digits = base_digits() if len(digits) > base: digits = digits[:base] s = str(s) if s == digits[0]: return 0 elif s.startswith('-'): return -base2int(s[1:], base=base, strip=strip, digits=digits) i = 0 for d in s: try: v = digits.index(d) except ValueError as e: if strip: continue e.args = ("invalid digit for base {base}: {s}".format(base=base, s=s),) raise i *= base i += v return i def digit_map(a=0, b=9, digits=None): """ create a map (dict) mapping individual digits to their numerical value. the symbols used for the digits can be provided, otherwise the default list set via base_digits() is used >>> digit_map(1, 3) == { '1': 1, '2': 2, '3': 3 } True """ if digits is None: digits = base_digits() return dict((digits[i], i) for i in irange(a, b)) # int2words implementation for lang='en' (English) _numbers = { 0: 'zero', 1: 'one', 2: 'two', 3: 'three', 4: 'four', 5: 'five', 6: 'six', 7: 'seven', 8: 'eight', 9: 'nine', 10: 'ten', 11: 'eleven', 12: 'twelve', 13: 'thirteen', 15: 'fifteen', 18: 'eighteen', } _tens = { 1: 'teen', 2: 'twenty', 3: 'thirty', 4: 'forty', 5: 'fifty', 6: 'sixty', 7: 'seventy', 8: 'eighty', 9: 'ninety' } def _int2words(n, scale='short', sep='', hyphen=' '): if n in _numbers: return _numbers[n] if n < 0: return 'minus ' + _int2words(-n, scale, sep, hyphen) if n < 20: return _numbers[n % 10] + _tens[1] if n < 100: (d, r) = divmod(n, 10) x = _tens[d] if r == 0: return x return x + hyphen + _numbers[r] if n < 1000: (d, r) = divmod(n, 100) x = _int2words(d, scale, sep, hyphen) + ' hundred' if r == 0: return x return x + ' and ' + _int2words(r, scale, sep, hyphen) if n < 1000000: (d, r) = divmod(n, 1000) x = _int2words(d, scale, sep, hyphen) + ' thousand' if r == 0: return x if r < 100: return x + ' and ' + _int2words(r, scale, sep, hyphen) return x + sep + ' ' + _int2words(r, scale, sep, hyphen) try: return __int2words(n, scale, sep, hyphen) except IndexError: raise ValueError(sprintf('Number too large (scale: {scale})')) # from http://en.wikipedia.org/wiki/Names_of_large_numbers _illions = [ 'm', 'b', 'tr', 'quadr', 'quint', 'sext', 'sept', 'oct', 'non', 'dec', 'undec', 'duodec', 'tredec', 'quattuordec', 'quindec', 'sexdec', 'septendec', 'octodec', 'novemdec', 'vigint', 'unvigint', 'duovigint', 'tresvigint', 'quattuorvigint', 'quinquavigint', 'sesvigint', 'septemvigint', 'octovigint', 'novemvigint', 'trigint', 'untrigint', 'duotrigint', 'trestrigint', 'quattuortrigint', 'quinquatrigint', 'sestrigint', 'septentrigint', 'octotrigint', 'noventrigint', 'quadragint', ] def __int2words(n, scale='short', sep='', hyphen=' '): """ convert a large integer (one million or greater) to a string representing the number in English, using short or long scale. >>> __int2words(10**12, scale='short') 'one trillion' >>> __int2words(10**12, scale='long') 'one billion' """ if scale == 'short': (g, p, k) = (3, 1000, 2) elif scale == 'long': (g, p, k) = (6, 1000000, 1) else: raise ValueError('Unsupported scale type: ' + scale) i = (len(str(n)) - 1) // g (d, r) = divmod(n, p**i) w = _illions[i - k] + 'illion' x = _int2words(d, scale, sep, hyphen) + ' ' + w if r == 0: return x if r < 100: return x + ' and ' + _int2words(r, scale, sep, hyphen) return x + sep + ' ' + _int2words(r, scale, sep, hyphen) def int2words(n, scale='short', sep='', hyphen=' ', lang='en'): """ convert an integer to a string representing the number (in English). scale - 'short' (for short scale numbers), or 'long' (for long scale numbers) sep - separator between groups hyphen - separator between tens and units lang - language (only 'en' (for English) is currently accepted) >>> int2words(1234) 'one thousand two hundred and thirty four' >>> int2words(-7) 'minus seven' >>> int2words(factorial(13)) 'six billion two hundred and twenty seven million twenty thousand eight hundred' >>> int2words(factorial(13), sep=',') 'six billion, two hundred and twenty seven million, twenty thousand, eight hundred' >>> int2words(factorial(13), sep=',', hyphen='-') 'six billion, two hundred and twenty-seven million, twenty thousand, eight hundred' >>> int2words(factorial(13), scale='long') 'six thousand two hundred and twenty seven million twenty thousand eight hundred' >>> sorted(irange(1, 10), key=int2words) [8, 5, 4, 9, 1, 7, 6, 10, 3, 2] """ if lang != 'en': raise ValueError(sprintf("int2words: lang='{lang}' not implemented")) return _int2words(int(n), scale, sep, hyphen) # convert an integer to BCD (binary coded decimal) # same as: nconcat(nsplit(n, base=10), base=16) def int2bcd(n, base=10, bits_per_digit=4): """ convert integer n into BCD (Binary Coded Decimal) the base and bits_per_integer can be specified (if desired) >>> int2bcd(123456) 1193046 >>> int2bcd(123456) == 0x123456 True """ s = 1 if n < 0: (s, n) = (-1, -n) r = k = 0 while True: (n, x) = divmod(n, base) r += (x << k) if n == 0: break k += bits_per_digit return s * r # convert a map to a string: "(a=1, b=2, c=3)" def map2str(m, sort=1, enc="()", sep=", ", arr="="): """ convert a map into a string (usually for output). the map may be a dict() type object or a collection of (key, value) pairs. >>> map2str(dict(a=1, b=2, c=3)) '(a=1, b=2, c=3)' >>> map2str(multiset("banana")) '(a=3, b=1, n=2)' >>> map2str(zip("abc", irange(1, 3))) '(a=1, b=2, c=3)' """ fn = (sorted if sort else identity) if isinstance(m, dict): # dict s = join((concat(k, arr, m[k]) for k in fn(m.keys())), sep=sep) else: # (k, v) pairs s = join((concat(k, arr, v) for (k, v) in fn(m)), sep=sep) return (enc[0] + s + enc[1] if enc else s) ############################################################################### # specialised classes: ############################################################################### # Delayed Evaluation # delayed evaluation (see also lazypy) class Delay(object): # set to force immediate evaluation immediate = 0 def __init__(self, fn, *args, **kw): """ create a delayed evaluation promise for fn(*args, **kw). Note that if you want the arguments themselves to be lazily evaluated you will need to use: Delay(lambda: fn(expr1, expr2, opt1=expr3, opt2=expr4)) rather than: Delay(fn, expr1, expr2, opt1=expr3, opt2=expr4) example: x = Delay(expensive, 1) x.evaluated --> False x.value --> returns expensive(1) x.evaluated --> True x.value --> returns expensive(1) again, without re-evaluating it x.reset() x.evaluated --> False x.value --> returns expensive(1), but re-evaluates it """ self.fn = fn self.args = args self.kw = kw self.evaluated = False if self.immediate: self.evaluate() def __repr__(self): return self.__class__.__name__ + '(value=' + (repr(self.value) if self.evaluated else '') + ')' def evaluate(self): self.value = self.fn(*(self.args), **(self.kw)) self.evaluated = True return self.value def reset(self): del self.value self.evaluated = False def __getattr__(self, key): if key == 'value': return self.evaluate() else: raise AttributeError() @classmethod def iter(self, *args): """ create an iterator that takes a sequence of Delay() objects (or callable objects with no arguments) and evaluates and returns each one as next() is called. i = Delay.iter(Delay(expensive, 1), Delay(expensive, 2), Delay(expensive, 3)) next(i) --> evaluates and returns expensive(1) next(i) --> evaluates and returns expensive(2) next(i) --> evaluates and returns expensive(3) """ # use x.value for delay objects, and just try and call everything else return ((x.value if type(x) is self else x()) for x in args) ############################################################################### # Value Accumulator class Accumulator(object): """ A value accumulator. >>> a = Accumulator(fn=max) >>> for x in (6, 5, 4, 7, 3, 1): a.accumulate(x) >>> a.value 7 >>> a = Accumulator() >>> for x in irange(1, 9): a.accumulate(x) >>> a.value 45 >>> fdiv(a.value, a.count) 5.0 """ def __init__(self, fn=operator.add, fn1=identity, value=None, data=None, collect=0, count=0): """ create an Accumulator. The accumulation function and initial value can be specified. """ self.fn = fn # used to accumulate self.fn1 = fn1 # used to set initial value self.value = value if collect and data is None: data = [] self.collect = collect self.data = data self.target = None self.count = count def __repr__(self): return self.__class__.__name__ + '(value=' + repr(self.value) + ', data=' + repr(self.data) + ', count=' + str(self.count) + ')' def accumulate(self, v=1): """ Accumulate a value. If the current value is None then this value replaces the current value. Otherwise it is combined with the current value using the accumulation function which is called as fn(, v). """ self.count += 1 self.value = (self.fn1(v) if self.value is None else self.fn(self.value, v)) def accumulate_data(self, v, data, target=None): """ Accumulate a value, and check the accumulated value against a target value, and if it matches record the data parameter. You can use this to record data where some function of the data is at an extremum value. If the 'collect' parameter was set during initialisation, then all values that hit current target are recorded in a list. Otherwise only the most recent value is recorded. """ if target is None: target = v self.accumulate(v) # have we hit the target? if self.value == target: if self.collect: # we need to collect all data values with specified target measure if target != self.target: # the target has changed, start a new list self.data = [data] self.target = target else: # if the target is unchanged, append the current data self.data.append(data) else: # otherwise, just record the data verbatim self.data = data def accumulate_from(self, s): """accumulate values from iterable object """ for v in s: self.accumulate(v) return self def accumulate_data_from(self, s, value=0, data=1): """ accumulate values and data from iterable object . , can be an index into elements from or a function to extract the appropriate value from an element. """ fn = lambda i: (lambda x: x[i]) if not callable(value): value = fn(value) if not callable(data): data = fn(data) for x in s: self.accumulate_data(value(x), data(x)) return self # rs = Accumulator.multi(fns=[min, max]) @classmethod def multi(cls, *args, **kw): return MultiAccumulator(*args, **kw) # combine multiple Accumulator objects into a MultiAccumulator # rs = Accumulator.combine(Accumulator(fn=min), Accumulator(fn=max), ...) @classmethod def combine(cls, *args): self = MultiAccumulator() self.multi = rs return self # multiple accumulators: e.g. MultiAccumulator(fns=[min, max]) class MultiAccumulator(object): def __init__(self, fns, *args, **kw): self.multi = list(Accumulator(fn, *args, **kw) for fn in fns) def __repr__(self): return self.__class__.__name__ + '(' + repr(self.multi) + ')' def perform(self, fn, *args, **kw): for x in self.multi: fn(x, *args, **kw) def accumulate(self, v): self.perform(Accumulator.accumulate, v) def accumulate_data(self, v, data): self.perform(Accumulator.accumulate_data, v, data) def accumulate_from(self, s): self.perform(Accumulator.accumulate_from, list(s)) def accumulate_data_from(self, s, value=0, data=1): self.perform(Accumulator.accumulate_data_from, list(s), value=value, data=data) def __getitem__(self, i): return self.multi[i] ############################################################################### # Routines for dealing with polynomials # represent polynomial a + bx + cx^2 + dx^3 + ... as: # # [a, b, c, d, ...] # # so the polynomial can be reconstructed as: # # sum(c * pow(x, i) for (i, x) in enumerate(poly)) # # make a polynomial from (exponent, coefficient) pairs # (we can use enumerate() to reverse the process) def poly_from_pairs(ps, p=None): if p is None: p = [] for (e, c) in ps: if c != 0: x = e + 1 - len(p) if x > 0: p.extend([0] * x) p[e] += c return poly_trim(p) poly_to_pairs = enumerate # remove extraneous zero coefficients def poly_trim(p): while len(p) > 1 and p[-1] == 0: p.pop() return p # multiply two polynomials def poly_mul(p, q): (np, nq) = (len(p), len(q)) if np < nq: (p, q) = (q, p) return poly_from_pairs( ((i + j, a * b) for (i, a) in enumerate(p) for (j, b) in enumerate(q)), [0] * (np + nq - 1) ) poly_zero = [0] poly_unit = [1] # multiply a sequence of polynomials def poly_multiply(ps): r = poly_unit for p in ps: r = poly_mul(r, p) return r # raise a polynomial to a (non-negative) integer power def poly_pow(p, n): n = as_int(n, include='0+') r = poly_unit while n > 0: (n, m) = divmod(n, 2) if m: r = poly_mul(r, p) if n: p = poly_mul(p, p) return r # add two polynomials def poly_add(p, q): return poly_from_pairs(enumerate(p), list(q)) # add a sequence of polynomials def poly_sum(ps): r = poly_zero for p in ps: r = poly_add(r, p) return r # map a function over the coefficients of a polynomial def poly_map(p, fn): return poly_trim(list(fn(x) for x in p)) # negate a polynomial def poly_neg(p): return list(-c for c in p) # subtract two polynomials def poly_sub(p, q): return poly_add(p, poly_neg(q)) # divide two polynomials # div() is used for coefficient division, if the leading coefficient of q is not 1 # (you probably want to use a rational implementation such as fractions.Fraction) def poly_divmod(p, q, div=rdiv): fn = (identity if q[-1] == 1 else (lambda x: div(x, q[-1]))) (d, r) = (poly_zero, p) while r != poly_zero: k = len(r) - len(q) if k < 0: break m = poly_from_pairs([(k, fn(r[-1]))]) d = poly_add(d, m) r = poly_sub(r, poly_mul(m, q)) return (d, r) # compose two polynomials: compose(p, q)(x) = p(q(x)) def poly_compose(p, q): r = poly_zero m = poly_unit for (i, a) in enumerate(p): if a: r = poly_add(r, list(a * c for c in m)) m = poly_mul(m, q) return r # print a polynomial in a more friendly form def poly_print(p, var='x'): r = list() for (e, c) in enumerate(p): if c == 0: continue s = str(c) if not (c < 0): s = '+' + s s = '(' + s + ')' if e == 0: pass elif e == 1: s = s + var else: s = s + var + '^' + str(e) r.append(s) return join(r[::-1], sep=" ") or "(0)" # evaluate a polynomial def poly_value(p, x): v = 0 for n in reversed(p): v *= x v += n return v # derivative of a polynomial def poly_derivative(p, k=1): for _ in irange(1, k): p = poly_from_pairs((e - 1, e * c) for (e, c) in enumerate(p) if e > 0) return p # integral of a polynomial (with constant c) def poly_integral(p, c=0, div=rdiv): k = c p = poly_from_pairs((e + 1, div(c, e + 1)) for (e, c) in enumerate(p)) p[0] = k return p # polynomial interpolation from a number of points def poly_interpolate(ps, field=None): ps = list(ps) k = len(ps) if k == 0: return None if k == 1: return [ps[0][1]] k -= 1 (A, B) = (list(), list()) for (x, y) in ps: A.append([1, x] + [x**i for i in irange(2, k)]) B.append(y) try: return poly_trim(Matrix.linear(A, B, field=field)) except ValueError: return # scale a polynomial to give integer coefficents def poly_scale(p, F=None): if not p: return p if F is None: F = Rational() p = list(map(F, p)) m = mlcm(*(f.denominator for f in p)) p = list(int(m * f) for f in p) g = mgcd(*p) if g > 1: p = list(x // g for x in p) return p # find rational roots of a polynomial # see: [ https://en.wikipedia.org/wiki/Rational_root_theorem ] def poly_rational_roots(p, domain="Q", include="+-0", F=None): """ find rational roots for the polynomial p (with rational coefficients). returns rational values x, such that: p(x) = 0 the type of roots returned can be controlled with the 'domain' and 'include' parameters: domain='Q' - find rational roots domain='Z' - find integer roots include='+' - include positive roots include='0' - include zero include='-' - include negative roots """ assert domain in "QZ" if not p: return (pos, neg, zero) = (x in include for x in '+-0') if F is None: F = Rational() # first deal with a root at x=0 if p[0] == 0: if zero: yield (0 if domain == "Z" else F(0, 1)) while p and p[0] == 0: p = p[1:] if not p: return # make an equivalent polynomial with integer coefficients p = poly_scale(p, F=F) fs = product(divisors(abs(p[0])), divisors(abs(p[-1]))) for x in uniq(map(unpack(F), fs)): if domain == "Z": if x.denominator != 1: continue x = x.numerator if pos and poly_value(p, x) == 0: yield x if neg and poly_value(p, -x) == 0: yield -x # return (n, r) where p = q^n . r def poly_div(p, q, div=rdiv): n = 0 while True: (p_, z) = poly_divmod(p, q, div=div) if z != poly_zero: break n += 1 p = p_ return (n, p) # EXPERIMENTAL: return factors of polynomial