#! 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: Sun Oct 11 21:02:59 2026 (Jim Randell) jim.randell@gmail.com # Language: Python (Python 2.7), Python3 (Python 3.6 - 3.15) # Package: N/A # Status: Free for non-commercial use # URI: http://www.magwag.plus.com/jim/enigma.html # # (c) Copyright 2009-2026, Jim Randell, all rights reserved. # ############################################################################### # -*- mode: Python; python-indent-offset: 2; coding: ascii -*- """ 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 - treat argument as an integer (where possible) base - get/set default base (radix) used in base conversion base_digits - get/set digits used in numerical base conversion base2int, str2int - convert a string in the specified base to an integer 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 cachegen - cache a generator (to allow indexed access) 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 a sequence in chunks clock - clock arithmetic variant on mod() clump, clump_val - 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 crt - solve a collection of modular congurences cslice - cumulative slices of an array csum - cumulative sum cycles - generate cycles of a permutation decompose - find sequences of integers with the specified sum diff - sequence difference digit_map - create a map of digits to corresponding integer values digrt - the (additive) digital root of a number disjoint_cproduct - disjoint cartesian product disjoint_union - disjoint union distinct_chars - check arguments (as strings) consist of distinct characters 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 exec_file - execute a file containing Python code express - express an amount using specific denominations factor - the prime factorisation of a number factorial - factorial function false - a function that always returns false farey - generate Farey sequences of coprime pairs fcompose - forward functional composition fdiv - float division fib - generate fibonacci sequences filter2 - partition a sequence into values that satisfy a predicate, and those that do not filter_unique - partition a sequence 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 a location where a function has a specified value find_values - find all locations where a function has a specified value find_zero - find a location where a function is zero first, ifirst - return items from the start of a sequence flatten, iflatten - flatten a list of lists flattened - fully flatten a nested structure floor - generalised floor function format_recurring - output the result from recurring() fraction, 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, ihypot - calculate hypotenuse icount, icount_* - count the number of elements of a sequence that satisfy a predicate identity - the identity function implies - logical implication (p -> q) int2base, int2str - 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 sequences 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 ipowers - generate perfect powers irange - inclusive range iterator irangef - inclusive range iterator with fractional steps iroot - integer kth root function is_coprime - check numbers are pairwise 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_ipower - check a number is a perfect power 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, nrev - 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 permutation - make a permutation from a set of cycles 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 prime_factors - generate primes factors of a number printf - print with interpolated variables pythagorean_triples - generate Pythagorean triples quadratic - find roots of a quadratic equation ratio, ratio_q - find lowest terms integer ratio rational - represet a rational number rcompose - reverse functional composition rcs, ircs - root of combined squares recurring - decimal representation of fractions recurring2fraction - find the fraction corrresponding to a decimal expansion repeat - generate repeated applications of a function repdigit - number consisting of repeated digits repeat - repeatedly apply a function to a value restrict - the restriction of a container to certain keys reverse, rev - 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 seq_get - get an item from a sequence seq_ordered - return an ordered tuple from the specified sequence singleton - return the value from a single valued container snapf - snap a float to an integer 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 sqrtmod - find modular square roots sqrx - perform a step in the extraction of a square root static - decorator to simulate static variables subfactorial - subfactorial function subsets, subseqs - generate subsequences of a sequence substitute - substitute symbols for digits in text substituted_expression - a substituted expression (alphametic/cryptarithm) solver substituted_sum - a solver for substituted sums sum_of_reciprocals - decompose a fraction into a sum of reciprocals (unit fractions) 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 true - a function that always returns True 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 a sequence unzip - inverse of zip unpack - return a function that unpacks its arguments update - create an updated version of a container object wrap - a decorator to wrap a function in another function zip_eq - check sequences contain the same elements 2D geometry functions: circle_intersect_circle - find where two circles intersect circle_intersect_line - find where a circle and line intersect circle_param - return a parametric form of a circle circle_param_t - return the parameter for a point on a circle line_bisect - find the perpendicular bisector of a line line_distance - minimum distance from a point to a line line_intersect - find the intersection of two lines line_param - return a parametric form of a line line_slope_intercept - return (slope, intercept) for a line point_distance - straight line distance between two points polygon_area - calculate the area of a simple polygon polygon_area_regular - calculate the area of a regular polygon triangle_circumcircle - find the circumcentre and circumradius of a triangle triangle_height - find the altitude of a triangle triangle_point - find the third vertex of a triangle 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 Enumerator - a class for iterating through a sequence and counting the items Football - a class for solving football league table puzzles MagicSquare - a class for solving magic squares Interval - an implementation of closed integer intervals Matrix - a class for manipulation 2d matrices multiset - an implementation of multisets (bags) Output - format blocks of output 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 """ # use Python 3 style print() and division in Python 2 from __future__ import (print_function, division) __author__ = "Jim Randell " __version__ = "2026-10-05" # -- __credits__ = "contributors = Brian Gladman; Frits ter Veen" import sys import os import operator import math import functools import itertools import collections import re # (see also the "six" module) _pythonv = sys.version_info[0:2] # Python version e.g. (2, 7) or (3, 11) if _pythonv[0] < 3: # Python 2.x _python = 2 if _pythonv < (2, 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 Mapping = collections.Mapping else: # 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, Mapping) else: Sequence = collections.Sequence Iterable = collections.Iterable Mapping = collections.Mapping _builtin_min = min _builtin_max = max _builtin_sum = sum # 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() module = sys.modules.get partial = functools.partial fsum = math.fsum # detect if running under PyPy _pypy = getattr(sys, 'pypy_version_info', None) # useful constants enigma = module(__name__) nl = "\n" pi = math.pi two_pi = pi + pi inf = float('+inf') fset = frozenset # "fset" is a shortcut for "frozenset" empty = fset() # the empty set str_upper = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" # upper case letters str_lower = "abcdefghijklmnopqrstuvwxyz" # lower case letters str_digit = "0123456789" # decimal digits str_nl = nl # PY_ENIGMA: # v = enable verbose output # p = enable prompts in arg() # i = enable interaction _PY_ENIGMA = os.getenv("PY_ENIGMA") or '' version = lambda: __version__ def import_fn(spec, verbose=0): """ import a value via a qualified spec, for example: 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 """ # we could use importlib.import_module() here importer = lambda x: __import__(x, fromlist=['']) if '.' not in spec: if verbose or 'v' in _PY_ENIGMA: print(str.format("[import_fn] importing \"{spec}\" ...", spec=spec)) return importer(spec) (mod, fn) = spec.rsplit('.', 1) if verbose or 'v' in _PY_ENIGMA: print(str.format("[import_fn] importing \"{mod}\" ...", mod=mod)) return getattr(importer(mod), fn) # execute some python code, and return a namespace with the result def exec_file(path, name=None, verbose=0): # if path is not a string, assume it is a (spec, file) pair if not isinstance(path, basestring): path = call(parsepath, path) # execute the file if verbose or 'v' in _PY_ENIGMA: print(str.format("[exec_file] executing \"{path}\" ...", path=path)) run_path = import_fn('runpy.run_path') ns = run_path(path, run_name='') return make_namespace(name, ns) # lazy importer: # [Python 3.15 has: "lazy import ..." and "lazy from X import Y", which can be the default: PYTHON_LAZY_IMPORTS=all] class LazyImporter(object): def __init__(self, spec, **kw): self.spec = spec self.name = (kw.get('name') or str.rsplit(spec, '.', 1)[-1]) self.space = (kw.get('space') or sys._getframe(1).f_globals) # replace in defining space self.verbose = (kw.get('verbose') or ('v' in _PY_ENIGMA)) self.obj = None # do the import def do_import(self): if self.verbose: printf("[LazyImporter] importing {self.spec}") obj = import_fn(self.spec) if self.name and self.space: # attempt to replace the name in the appropriate namespace try: self.space[self.name] = obj except Exception as e: pass # cache the object self.obj = obj return obj # function call def __call__(self, *args, **kw): #print(['__call__', self, args, kw]) # do the import fn = self.obj or self.do_import() # call the real function return fn(*args, **kw) # attribute fetch def __getattr__(self, key): #print(['__getattr__', self, key]) if key == '__wrapped__': return # otherwise doctest fails # do the import obj = self.obj or self.do_import() # return the real attribute return getattr(obj, key) lazy_import = LazyImporter # 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). # TODO: add a parameter for adding a 'key' function 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+). """ cache = dict() try: f.cache = cache except AttributeError: pass @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) # cache a generator def cachegen(seq): """ allow indexed access to a generator. the returned object is a callable that can be passed a (non-negative) index as argument to return the value at that index. values are cached in the 'cache' attribute of the returned object. >>> f = cachegen(fib(0, 1)) >>> f(30) 832040 >>> f(10) 55 """ seq = iter(seq) cache = list() # get the value at index k def fn(k, seq=seq, cache=cache): if k < 0: raise IndexError(str.format("invalid index {k}", k=k)) n = k - len(cache) if n >= 0: cache.extend(first(seq, n + 1)) return cache[k] fn.cache = cache return fn # 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 a false value def false(*args, **kw): # type: (...) -> bool """a function that ignores any arguments and returns False""" return False # 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 except Exception: pass if 'default' in kw: return kw['default'] msg = "invalid integer: " + repr(x) if include: msg += ' [include: ' + include + ']' raise ValueError(msg) #is_int = partial(as_int, default=None) # division that always produces a float 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) # less than/greater than (or equal) to a target; useful for filter() etc. def eq(t): "return a function f(x) = (x == t)"; return (lambda x: x == t) def lt(t): "return a function f(x) = (x < t)"; return (lambda x: x < t) def le(t): "return a function f(x) = (x <= t)"; return (lambda x: x <= t) def gt(t): "return a function f(x) = (x > t)"; return (lambda x: x > t) def ge(t): "return a function f(x) = (x >= t)"; return (lambda x: x >= t) def between(a, b): "return a function f(x) = (a < x < b)"; return (lambda x: a < x < b) # exclusive between def betweene(a, b): "return a function f(x) = (a <= x <= b)"; return (lambda x: a <= x <= b) # inclusive between # membership/non-membership of a collection def is_in(s): "return a function f(x) = (x in s)"; return (lambda x: x in s) def is_not_in(s): "return a function f(x) = (x not in s)"; return (lambda x: x not in s) def inc(i=1): """return a function that increments by a fixed amount""" 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. def compare(a, b, vs=None): """ return -1 if a < b, 0 if a == b and +1 if a > b. can be specified as a triple of values to return instead of (-1, 0, +1) respectively. >>> 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) # negation neg = operator.neg # logical implication: p -> q, "if p then 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 q or (not p) # 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 seq_is_distinct(seq, n=None): """ returns True if the finite sequence has no repeated values, or if is specified contains exactly different values. alias: distinct_values(). """ seq = list(seq) if n is None: n = len(seq) return len(set(seq)) == n distinct_values = seq_is_distinct # this is the same as: seq_multiplicity(seq, n=1) def seq_all_different(*seqs, **kw): """ check all elements of are pairwise distinct. if multiple sequences are provided elements must be distinct across all sequences. the elements of the sequences must be hashable. >>> 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 """ fn = kw.pop('fn', None) if kw: raise TypeError(str.format("seq_all_different: unknown arguments {kw}", kw=seq2str(kw.keys()))) seen = set() for seq in seqs: for x in seq: if fn: x = fn(x) if x in seen: return False seen.add(x) return True # generalised version of seq_all_different() def seq_multiplicity(*seqs, **kw): """ check all elements of have multiplicity . if multiple sequences are provided then the total multiplicity across all sequences is considered. >>> seq_multiplicity([0, 1, 2, 3], n=1) True >>> seq_multiplicity([2, 1, 2, 3], [0, 1, 0, 3], n=2) True """ n = kw.pop('n', 1) fn = kw.pop('fn', None) if kw: raise TypeError(str.format("seq_multiplicity: unknown arguments {kw}", kw=seq2str(kw.keys()))) seen = multiset() for seq in seqs: for x in seq: if fn: x = fn(x) if seen.get(x, 0) == n: return False seen.add(x) return seen.all_same(n) # 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, **kw) 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. """ itr = iter(seq) n = 0 # is value specified? if 'value' in kw: v = kw['value'] else: # otherwise, use the first value try: v = next(itr) n = 1 except StopIteration: # empty sequence return Record(same=True, empty=True, value=None) # check the rest of the sequence for x in itr: 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.irange(1, inf)) 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 is 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 """ reverse = kw.pop('reverse', 0) strict = kw.pop('strict', None) k = kw.pop('first', None) if kw: raise TypeError(str.format("zip_eq: unknown arguments {kw}", kw=seq2str(kw.keys()))) if reverse: ss = (reversed(s) for s in ss) z = (zip(*ss, strict=strict) if strict is not None else zip(*ss)) if k is not None: z = first(z, count=k, fn=iter) # was: [[ return all(seq_all_same(vs) for vs in z) ]] for vs in z: if not seq_all_same(vs): return False return True # NOTE: not sure these will remain in this form # find positions where sequences are correlated def zip_corr(*ss, **kw): """ return indices where sequences have the same element values. the 'strict' argument is passsed to zip (which is 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 """ strict = kw.pop('strict', None) k = kw.pop('first', None) if kw: raise TypeError(str.format("zip_correlate: unknown arguments {kw}", kw=seq2str(kw.keys()))) z = (zip(*ss, strict=strict) if strict is not None else zip(*ss)) if k is not None: z = first(z, count=k, fn=iter) for (i, vs) in enumerate(z): if seq_all_same(vs): yield i # just count the number of correlations zip_ncorr = lambda *ss, **kw: icount(zip_corr(*ss, **kw)) # special case the most usual scenario = count the correlations between two lists zip_ncorr2 = lambda xs, ys: sum(1 for (x, y) in zip(xs, ys) if x == y) def ordered(*args, **kw): """ return 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,) """ fn = kw.pop('fn', tuple) return fn(sorted(args, **kw)) def seq_ordered(seq, **kw): """ return as a tuple in order >>> seq_ordered([2, 1, 3]) (1, 2, 3) >>> seq_ordered([2, 1, 3], reverse=1) (3, 2, 1) """ fn = kw.pop('fn', tuple) return fn(sorted(seq, **kw)) # is sequence sorted? def is_sorted(seq, strict=0, fn=operator.lt): """ check if is sorted. adjacent values must satisfy (default is: <) if is not set, then adjacent values may also be equal """ # was: [[ return all(f(x, y) for (x, y) in tuples(seq, 2)) ]] if strict: for (x, y) in tuples(seq, 2): if not fn(x, y): return False else: for (x, y) in tuples(seq, 2): if not (x == y or fn(x, y)): return False return True def is_increasing(seq, strict=0): "check if is increasing"; return is_sorted(seq, strict=strict, fn=operator.lt) def is_decreasing(seq, strict=0): "check if is decreasing"; return is_sorted(seq, strict=strict, fn=operator.gt) def is_consecutive(seq, gap=1, strict=1): "check is consist of consecutive values"; return is_sorted(seq, strict=strict, fn=(lambda x, y: x + gap == y)) def encl(s, b="{}", fn=str): """ enclose string by bracketing it with the elements of . is initially processed by . >>> encl('xyz') '{xyz}' >>> encl(42, b='[]') '[42]' >>> encl('xyz', '|') '|xyz|' """ return b[0] + fn(s) + b[-1] # create a function to apply a specific format spec def fmts(spec, *rest): """ fmts(spec, object) = format an object using the given format specifier. fmts(spec) = returns a function that takes an object and formats the object using the given format specifier. this can be useful in the 'fn' parameter of join()/concat()/seq2str(). >>> fmts(".4f", pi) '3.1416' >>> fmts("pi = {.4f}", pi) 'pi = 3.1416' >>> fmts("pi = {.4f}")(pi) 'pi = 3.1416' """ if '{' not in spec: # standard spec if not rest: return (lambda obj: format(obj, spec)) if len(rest) == 1: return format(rest[0], spec) else: # embedded spec if not rest: return (lambda obj: fmts(spec, obj)) if len(rest) == 1: obj = rest[0] return re.sub(r'{(.*?)}', (lambda x: format(obj, x.group(1))), spec) # I would prefer join() to be a string constructor: # str.from_seq(seq, sep='', enc=''), or just: str.join(seq, sep='', enc='') # but here is a utility function (aka: seq2str) def join(seq, sep='', enc='', sort=0, rev=0, 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' """ if sort: seq = sorted(seq) if rev: seq = reversed(seq) r = str.join(sep, (fn(x) for x in seq)) if enc: r = str.join('', (enc[0], r, enc[-1])) return r def joinf(sep='', enc='', sort=0, rev=0, fn=str): "return a joining function" return partial(join, sep=sep, enc=enc, sort=sort, rev=rev, 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.pop('sep', '') enc = kw.pop('enc', '') sort = kw.pop('sort', 0) rev = kw.pop('rev', 0) fn = kw.pop('fn', str) if kw: raise TypeError(str.format("concat: unknown arguments {kw}", kw=seq2str(kw.keys()))) 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, sort=sort, rev=rev, fn=fn) # reverse a sequence or a map (maybe -> rev()) # for a unicode string if will be considered as a sequence of codepoints # you might want to convert it to a sequence of graphemes instead def reverse(seq, 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. alias: rev(). >>> 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(seq, 'rev'): return (seq.rev() if fn is None else fn(seq.rev())) # if it is a dict, return a reverse map if isinstance(seq, Mapping): return type(seq)((v, k) for (k, v) in seq.items()) # if it is not already a sequence, turn it into one if not isinstance(seq, Sequence): seq = list(seq) # if it is a string, return a string if fn is None: if isinstance(seq, basestring): fn = join elif isinstance(seq, tuple): fn = tuple else: fn = list return fn(reversed(seq)) 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, Mapping): 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) # the default base for positional number systems (can be set via base()) # used if 'base' parameter is not specified (or is specified as 0) radix = 10 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 Python 3 [[ def nconcat(*digits, base=10): ]] is allowed instead base = kw.pop('base', 0) if base == 0: base = radix if kw: raise TypeError(str.format("nconcat: unknown arguments {kw}", kw=seq2str(kw.keys()))) # allow a sequence to be 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) def nsplitter(n, k=None, base=0, validate=0): """split integer into digits, starting with the least significant digit""" if base == 0: base = radix if base < 2: raise ValueError(str.format("invalid base: {base!r}", base=base)) if n == inf or n == -inf: raise ValueError(str.format("invalid number: {n!r}", n=n)) if validate: (n, base) = (as_int(n), as_int(base, include='+')) n = abs(n) 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=0, fn=tuple, reverse=0, validate=0): """ 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. if is set to a true value then the digits are presented in reverse order. >>> 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(sq(111), 3) (3, 2, 1) """ ds = nsplitter(n, k=k, base=base, validate=validate) if not reverse: ds = rev(ds) return fn(ds) def dsum(n, k=None, base=0, validate=0): """ 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, validate=validate)) # population count, Hamming weight, bitsum(), bit_count() dsum2 = getattr(int, 'bit_count', None) if not dsum2: 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=0, validate=0): """ return the number of digits in a number, when represented in the specified base. if the number is not an integer, only the integer part is considered. >>> ndigits(factorial(70)) 101 """ if n == inf or n == -inf: return inf return icount(nsplitter(n, base=base, validate=validate)) # TODO: maybe: ndigit(n, i, j, k, ...) -> extract digits i, j, k from n def is_ndigits(n, k, base=0, validate=0): """check if has digits in the specified base""" if base == 0: base = radix if n == inf or n == -inf: return (k == inf) if validate: n = as_int(n) n = abs(int(n)) x = base**(k - 1) if n < x: return False return n < x * base # return the smallest and largest k-digit numbers def ndigit_range(k, base=0, adj=None): """ return the smallest and largest k-digit numbers (in base ). the endpoints are calculated as base**(k - 1) and (base**k) - 1. these can be adjusted by passing a pair of values to . (e.g. pass adj=(0, 1) for to adjust the values to be suitable for Python's builtin range() function). >>> ndigit_range(4) (1000, 9999) >>> ndigit_range(4, base=16) (4096, 65535) """ if base == 0: base = radix (lo, hi) = (base**(k - 1), (base**k) - 1) if adj: lo += adj[0] hi += adj[-1] return (lo, hi) # maybe -> nrev() def nreverse(n, k=None, base=0, validate=0): """ 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, validate=validate) return nconcat(nsplitter(n, k=k, base=base, validate=validate), base=base) nrev = nreverse fnmatch = lazy_import("fnmatch.fnmatchcase", name='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.fnmatchcase). 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=0): """ 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 numbers(ss, base=0, csep=',', crange='-', cneg='!', strip=1, enc='', fn=list): """ generate numbers (non-negative integers) according to the clauses specified in . clauses may be: ",...," = a sequence of clauses "-" = a range of numbers "" = a single number "!" = a clause to exclude clause separator is any character in . range and negation indicators are specified by and (which can be empty to disable operation) if strip=1 is specified non-digit characters are removed when parsing numbers. if enc is specified the specification should be enclosed. >>> numbers("1-9") [1, 2, 3, 4, 5, 6, 7, 8, 9] >>> numbers("1-3,7-9") [1, 2, 3, 7, 8, 9] >>> numbers("1-9,!4,!5,!6") [1, 2, 3, 7, 8, 9] >>> numbers("1-9,!4-6") [1, 2, 3, 7, 8, 9] >>> numbers("[1-9,!4-6]", enc="[]") [1, 2, 3, 7, 8, 9] """ ss = ss.strip() if enc: if not (ss[0] == enc[0] and ss[-1] == enc[-1]): raise ValueError("numbers: unenclosed spec " + repr(ss)) ss = ss[1:-1] # split into terms, separated by any character in d = csep[0] if len(csep) > 1: # map all other separators to d s = re.sub(encl(re.escape(csep[1:]), '[]'), d, ss) # collect positive and negative terms (tpos, tneg) = (list(), list()) for t in ss.split(d): if not t: raise ValueError("numbers: empty term") # indicates a negative term if cneg and t[0] == cneg: ts = tneg t = t[1:] else: ts = tpos # indicates a range if crange and crange in t: (a, _, b) = t.partition(crange) ts.extend(irange(base2int(a, base=base, strip=1), base2int(b, base=base, strip=strip))) continue # otherwise, a numeric literal ts.append(base2int(t, base=base, strip=strip)) # generate the numbers return diff(tpos, tneg, fn=fn) def respace(x, sep=' '): """ respace string . leading and trailing space is removed, and other space is reduced to a single space character. >>> respace(" hello world! ") 'hello world!' """ if not isinstance(x, basestring): return x # leave non-strings alone return str.join(sep, str.split(x)) 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(seq, 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 (seq if k == 0 else seq[k:] + seq[: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 in chunks of size . (for overlapping tuples see tuples()) see also: itertools.batched() (Python 3.12+) >>> 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)] >>> list(unzip(chunk(irange(0, 9), 2))) [(0, 2, 4, 6, 8), (1, 3, 5, 7, 9)] """ itr = iter(seq) while True: s = fn(itertools.islice(itr, 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_val(seq, fn=None): """ separate into contiguous blocks of repeated values (according to function ). returns (, ) pairs use clump() if you just want the values. see also: itertools.groupby(seq, fn). """ seq = iter(seq) try: x = next(seq) except StopIteration: return v = (x if fn is None else fn(x)) xs = [x] for x in seq: v_ = (x if fn is None else fn(x)) if v_ == v: xs.append(x) else: yield (v, xs) v = v_ xs = [x] yield (v, xs) def clump(seq, fn=None): """ separate into contiguous blocks of repeated values (according to function ). use clump_val() if you also want the common value. >>> 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]] """ for (v, xs) in clump_val(seq, fn): yield xs # set union of a bunch of sequences def union(ss, fn=set, fnu=None): """construct a set that is the union of the sequences in """ # formerly this was: [[ fn().union(*ss) ]] but that doesn't seem to work with generators if fnu is None: fnu = fn.update rs = fn() for s in ss: fnu(rs, s) return rs # 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: xs = fn(xs) if rs is None: rs = xs else: if not rs.isdisjoint(xs): return None rs.update(xs) return rs # check sequences in are disjoint is_disjoint = lambda ss, fn=set: disjoint_union(ss, fn=fn) is not None # ss = remaining sets to process # vs = selected values _dc_key_fn = (lambda k_v: len(k_v[1])) # (k, v) -> len(v) def _disjoint_cproduct(ss, vs): # final slot? if len(ss) == 1: [(j, xs)] = ss.items() for x in xs: vs[j] = x yield tuple(vs) else: # choose a slot with the fewest options (j, xs) = min(ss.items(), key=_dc_key_fn) for x in xs: ss_ = dict((k, v.difference([x])) for (k, v) in ss.items() if k != j) vs[j] = x #yield from _disjoint_cproduct(ss_, vs) #[Python 3] for z in _disjoint_cproduct(ss_, vs): yield z #[Python 2] # disjoint cartestian product of a sequence of sets, any value may only appear in one position def disjoint_cproduct(ss): d = dict() for (k, xs) in enumerate(ss): xs = set(xs) if not xs: return [] d[k] = xs if not d: return [()] return _disjoint_cproduct(d, [None] * len(d)) distinct_cproduct = disjoint_cproduct # set intersection of a bunch of sequences def intersect(ss, fn=set, **kw): """construct a set that is the intersection of the sequences in """ itr = iter(ss) try: s = fn(next(itr)) except StopIteration: pass else: return s.intersection(*itr) if 'default' in kw: return kw['default'] raise ValueError("empty intersection") # return an element of a container def peek(seq, k=0, p=None, **kw): """ return an element of a container. if k is specified the first k values chosen are discarded (so, for a sequence, you will get seq[k]). empty containers (or those with too few elements) return the specified 'default' value, or raise an IndexError. if p is specified, then only items satisfying predicate p will be considered. (i.e. other elements will be discarded). 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 kw.pop('validate', None): k = as_int(k, include="0+") if kw and list(kw.keys()) != ['default']: raise TypeError(str.format("peek: unknown arguments {kw}", kw=seq2str(kw.keys()))) if p is not None: seq = filter(p, seq) if k >= 0: if not isinstance(seq, Mapping): # try to index into the container try: return seq[k] except (KeyError, IndexError, TypeError): pass # try iterating through the container for (i, x) in enumerate(seq): if i == k: return x # return any default value try: return kw['default'] except KeyError: pass raise IndexError(str.format("invalid index {k}", k=k)) def seq_get(seq, k=None, default=None): """ get the item at index from sequence , or return the default value if the index is not valid (negative indices are not allowed). (use peek(seq, k) if you want to get an IndexError exception). if k is None, then a function is returned to get an item from the specified sequence . >>> seq_get([1, 2, 3], 1) 2 >>> seq_get([1, 2, 3], 4) is None True >>> seq_get("abc", -1) is None True >>> seq_get("abc")(1) 'b' """ if k is None: return partial(peek, seq, default=default) else: return peek(seq, k, default=default) def seq_items(seq, i=0, j=None): """ generate (, ) pairs from sequence . items are generated starting from index (i.e. the first items are skipped), and finish at index (if specified, otherwise to the end of the sequence). similar to: ((k, seq[k]) for k in irange(i, j)) >>> list(seq_items("ABCDEFGHI", 3)) [(3, 'D'), (4, 'E'), (5, 'F'), (6, 'G'), (7, 'H'), (8, 'I')] >>> list(seq_items("ABCDEFGHI", 3, 5)) [(3, 'D'), (4, 'E'), (5, 'F')] """ if i > 0: j = (None if j is None or j == inf else j + 1) return enumerate(itertools.islice(seq, i, j), start=i) if i == 0: if not (j is None or j == inf): seq = itertools.islice(seq, j + 1) return enumerate(seq) raise ValueError("seq_items: invalid start value: i={i}".format(i=i)) # 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): multi = kw.pop('multi', 0) if kw: raise TypeError(str.format("item: unknown arguments {kw}", kw=seq2str(kw.keys()))) f = operator.itemgetter(*ks) if len(ks) == 1 and multi: return (lambda x: (f(x),)) return f def attr(*ks, **kw): multi = kw.pop('multi', 0) if kw: raise TypeError(str.format("attr: unknown arguments {kw}", kw=seq2str(kw.keys()))) f = operator.attrgetter(*ks) if len(ks) == 1 and multi: return (lambda x: (f(x),)) return f items = lambda n: map(item, xrange(n)) # e.g. (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) # make a range from a slice on a sequence of length n def make_range(i, n): (a, b, c) = (i.start, i.stop, i.step) a = (0 if a is None else (a + n if a < 0 else a)) b = (n if b is None else (b + n if b < 0 else b)) c = (1 if c is None else c) return range(a, b, c) # a circular list, where indices are calculated mod length class CircularList(list): def __getitem__(self, i): n = len(self) # deal with a slice argument -> return a list if isinstance(i, slice): return list(self[j] for j in make_range(i, n)) # otherwise, integer index -> return an element if n == 0: raise IndexError("list index out of range") return list.__getitem__(self, operator.index(i) % n) # list.copy() returns a list(), so re-wrap it def copy(self): return self.__class__(list.copy(self)) def diff(a, b, *rest, **kw): """ return the subsequence of that excludes elements in (and ). the result is returned as a tuple, or constructed using . alias: excl(). >>> diff((1, 2, 3, 4, 5), (3, 5, 2)) (1, 4) >>> diff('newhampshire', 'wham', fn=join) 'nepsire' """ fn = kw.pop('fn', tuple) if kw: raise TypeError(str.format("diff: unknown arguments {kw}", kw=seq2str(kw.keys()))) b = set(b) if rest: b.update(*rest) return fn(x for x in a if x not in b) # or: excl(ude) / incl(ude) excl = diff def incl(a, b, *rest, **kw): """ return the subsequence of that includes elements in . >>> incl((1, 2, 3, 4, 5), (2, 3, 5, 7)) (2, 3, 5) """ fn = kw.pop('fn', tuple) if kw: raise TypeError(str.format("incl: unknown arguments {kw}", kw=seq2str(kw.keys()))) if rest: b = set(b).union(*rest) return fn(x for x in a if x 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 (sometimes known as: bags) 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) if you call freeze() on a multiset it becomes hashable, but once a multiset is frozen you should not change it (although this is not enforced, unless you set [[ multiset.enforce_freeze = 1 ]]). """ # set for enforced frozen multisets (0 = faster; 1 = enforced immutability) # (useful during development to catch "accidental" modifications) enforce_freeze = 0 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) self.frozen = 0 self._hash = None # deal with any initialisation objects for v in vs: if isinstance(v, Mapping): # from a dict self.update_from_dict(v, validate=1) else: # from a sequence v = list(v) try: s = multiset().update_from_pairs(v, validate=1) self.update_from_dict(s, validate=1) except (TypeError, ValueError): # maybe more, or maybe just Error self.update_from_seq(v, validate=1) # add in any keyword items if kw: for (x, n) in kw.items(): self.add(x, count=n, validate=1) def update_from_seq(self, vs, count=1, validate=0): """update a multiset from a sequence of items""" for x in vs: self.add(x, count=count, validate=validate) return self def update_from_pairs(self, vs, validate=0): """update a multiset from a sequence of (, ) pairs""" for (x, n) in vs: self.add(x, count=n, validate=validate) return self def update_from_dict(self, d, validate=0): """update a multiset from a dict of -> values""" return self.update_from_pairs(d.items(), validate=validate) @classmethod def from_dict(cls, *vs, **kw): """ create a multiset from a dict of -> values (or multiple dicts). """ validate = kw.pop('validate', 0) if kw: raise TypeError(str.format("multiset.from_dict: unknown arguments {kw}", kw=seq2str(kw.keys()))) m = cls() for v in vs: m.update_from_dict(v, validate=validate) return m @classmethod def from_pairs(cls, *vs, **kw): """ create a multiset from a sequence of (, ) pairs (or multiple sequences). """ validate = kw.pop('validate', 0) if kw: raise TypeError(str.format("multiset.from_pairs: unknown arguments {kw}", kw=seq2str(kw.keys()))) m = cls() for v in vs: m.update_from_pairs(v, validate=validate) return m @classmethod def from_seq(cls, *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.pop('count', 1) validate = kw.pop('validate', 0) if kw: raise TypeError(str.format("multiset.from_seq: unknown arguments {kw}", kw=seq2str(kw.keys()))) m = cls() for v in vs: m.update_from_seq(v, count=count, validate=validate) 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) def is_duplicate(self, n=1): "does this multiset contain elements with multiplicity greater than ?" # was: [[ return any(v > n for v in dict.values(self)) ]] for v in dict.values(self): if v > n: return True return False def all_same(self, n=None): "does this multiset contain only values with the same multiplicity ( if specified)" if n is None: return seq_all_same(dict.values(self)) # was: [[ return all(v == n for v in dict.values(self)) ]] for v in dict.values(self): if v != n: return False return True # 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): return len(dict.keys(self)) # specialised version for PyPy3 def __distinct_size_pypy3(self): return len(dict(self)) distinct_size = (__distinct_size_pypy3 if (_pypy and _python > 2) else __distinct_size) distinct_size.__doc__ = "return the number of distinct elements in the multiset" # return the count of an item def count(self, item): """return the number of times occurs in the multiset""" return dict.get(self, item, 0) # freeze a multiset def freeze(self): """ freeze a multiset. after it is frozen it may be hashed, which allows it to be inserted into a set or used as the key of a dict. but once frozen the multiset must not be changed. (to help spot accidental attempts to change a multiset after freezing set [[ multiset.enforce_freeze = 1 ]]. multiset.copy() can be used to create an unfrozen copy of a frozen multiset. """ # if we are being strict we return a new frozen multiset that cannot be modified if multiset.enforce_freeze: return fmultiset(self) # otherwise, we just mark this as frozen (which enables hashing) # (if an equivalent dict.freeze() existed, we could call it here) self.frozen = 1 return self # hash a frozen multiset def __hash__(self): if not self.frozen: raise TypeError("attempt to hash an unfrozen multiset") if self._hash is None: self._hash = hash(fset(self.to_pairs())) return self._hash # add an item def add(self, item, count=1, validate=0): """ add an item to a multiset. count can be negative to remove items. """ if validate: count = as_int(count) try: count += self[item] if count == 0: del self[item] return self except KeyError: pass if count < 0: raise ValueError(str.format("multiset.add: negative count for {item!r} -> {count}", item=item, count=count)) if count > 0: self[item] = count return self # remove an item def remove(self, item, count=1, validate=0): """remove an item from the multiset""" return self.add(item, count=-count, validate=validate) # delete an item (no error is raised if the item does not exist) def delete(self, item): """delete all occurrences of an item from the multiset""" return self.pop(item, 0) discard = delete # 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 is specifed only the first items are returned. """ s = sorted(dict.items(self), key=(lambda t: t[::-1]), reverse=1) return (s if n is None else s[:n]) def multiplicity(self, n): """ return elements with multiplicity may be specified as a callable, for example: eq(3), lt(4). """ fn = (n if callable(n) else eq(n)) for (k, v) in dict.items(self): if fn(v): yield (k, v) # provide some useful operations on multisets # update self with some other multisets (item counts are summed) def update(self, *rest, **kw): """ update the multiset with some other multisets (or objects that can be interpreted as multisets). item counts are summed. """ validate = kw.pop('validate', 0) if kw: raise TypeError(str.format("multiset.update: unknown arguments {kw}", kw=seq2str(kw.keys()))) for m in rest: if not isinstance(m, Mapping): m = multiset(m) self.update_from_dict(m, validate=validate) return self # combine self and some other multisets (item counts are summed) def combine(self, *rest, **kw): """ 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. """ validate = kw.pop('validate', 0) if kw: raise TypeError(str.format("multiset.combine: unknown arguments {kw}", kw=seq2str(kw.keys()))) return multiset(self).update(*rest, validate=validate) # 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, Mapping): 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 self.copy().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 = self.copy() for m in rest: if not isinstance(m, Mapping): m = multiset(m) r = multiset.from_pairs((item, min(count, r.get(item, 0))) for (item, count) in m.items()) return r # is this multiset a subset of m? def issubset(self, m, strict=0): """test if the multiset is contained in multiset """ if not isinstance(m, Mapping): m = multiset(m) # check sizes r = compare(self.size(), m.size()) if r == 1 or (strict and r == 0): return False # check items # was: [[ return all(count <= m.get(item, 0) for (item, count) in self.items()) ]] for (item, count) in self.items(): if count > m.get(item, 0): return False return True # is this multiset m a superset of self? def issuperset(self, m, strict=0): """test if the multiset contains multiset """ if not isinstance(m, Mapping): m = multiset(m) return m.issubset(self, strict=strict) # 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, Mapping): 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=count) elif count < 0: d2.add(item, count=-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] # 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, Mapping): m = multiset(m) if any(x in self for x in m): return False return True isdisjoint = is_disjoint # 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 dict.items(self)) # how many subsets of this multiset def powerset_size(self): return multiply(v + 1 for v in dict.values(self)) 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) es = list(self.elements()) for k in irange(min_size, max_size): for ss in uC(es, k): yield multiset.from_seq(ss) def copy(self): """return a copy of the multiset""" return multiset.from_dict(self) def min(self, **kw): """ return the minimum item value in a multiset (or ). equivalent to: min(self) """ if (not self) and 'default' in kw: return kw['default'] if 'key' in kw: return _builtin_min(dict.keys(self), key=kw['key']) return _builtin_min(dict.keys(self)) def max(self, **kw): """ return the maximum item value in a multiset (or ). equivalent to: max(self) """ if (not self) and 'default' in kw: return kw['default'] if 'key' in kw: return _builtin_max(dict.keys(self), key=kw['key']) return _builtin_max(dict.keys(self)) def sum(self, fn=_builtin_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 avg(self, div=fdiv, fn=_builtin_sum): """ return the average (arithmetic mean) of the items in a multiset. """ return div(self.sum(fn=fn), self.size()) def multinomial(self): """ The number of different arrangements of this multiset. >>> multiset("banana").multinomial() 60 """ return multinomial(self.values()) def mpermutations(self, k=None): """ return all possible arrangements of this multiset. >>> icount(multiset("banana").mpermutations()) 60 """ return mpermutations(self, k=k) def express(self, v, k=None, fn=identity): """ express value using a subset of this multiset. if is specified, only find subsets of size . return subsets of the original multiset, optionally processed by . """ ps = sorted(dict.items(self)) for rs in express_pairs(v, ps, tv=self.sum(), k=k): yield fn(multiset.from_pairs(rs)) def expressible(self, vs, k=None): """ are all values in expressible using subsets of this multiset? see: multiset.express() """ ps = sorted(dict.items(self)) t = self.sum() # was: [[ return all(any(express_pairs(v, ps, t, k)) for v in vs) ]] for v in vs: if not any(express_pairs(v, ps, tv=t, k=k)): return False return True def sorted(self, key=None, reverse=False): """generate elements of the multiset in order""" for k in sorted(dict.keys(self), key=key, reverse=reverse): for _ in xrange(self.get(k)): yield k # restriction of a multiset to a specific set of keys def restrict(self, ks, strict=0): return restrict(self, ks, strict=strict) # find exact covers of this multiset def mcover(self, m, reject=None): return mcover(m, self, reject=reject) def map2str(self, sort=1, enc='()', sep=', ', arr='='): """call map2str() on the multiset""" return map2str(self, sort=sort, enc=enc, sep=sep, arr=arr) def fmt(self, flags='d', sort=1, rev=0): """ format a multiset as a string. the first flag controls the formatting: d = dictionary formatting l = list formatting p = pairs formatting additional flags: r = use repr() on elements c = include constructor function """ r = "???" f = flags[0] if f in 'dp': ks = self.keys() if sort: ks = sorted(ks, reverse=rev) fn = (repr if 'r' in flags else identity) if f == 'd': r = map2str(((fn(k), self.get(k)) for k in ks), arr=": ", sep=", ", enc="{}") if 'c' in flags: r = "multiset.from_dict({r})".format(r=r) elif f == 'p': ps = ("({k}, {v})".format(k=fn(k), v=self.get(k)) for k in ks) r = seq2str(ps, sep=", ", enc="[]") if 'c' in flags: r = "multiset.from_pairs({r})".format(r=r) elif f == 'l': # format as a list seq = self if sort: seq = seq.sorted(reverse=rev) if 'r' in flags: seq = map(repr, seq) r = seq2str(seq, sep=", ", enc="[]") if 'c' in flags: r = "multiset.from_seq({r})".format(r=r) return r # generate item pairs def to_pairs(self): return tuple(dict.items(self)) 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 # ^ = symmetric_difference # * = multiply # < = subset # > = superset __add__ = combine __sub__ = difference __and__ = intersection __or__ = union __iadd__ = update __ior__ = union_update __xor__ = symmetric_difference __mul__ = multiply __le__ = issubset __ge__ = issuperset __lt__ = lambda self, m: self.issubset(m, strict=1) __gt__ = lambda self, m: self.issuperset(m, strict=1) # enforced frozen multiset class fmultiset(multiset): # can only be initialised from an existing multiset def __init__(self, m=None): multiset.__init__(self) if m: dict.update(self, m) self.frozen = 1 def _frozen_error(self, *args, **kw): raise TypeError("attempt to modify frozen multiset") # override multiset methods that might change the multiset add = delete = union_update = _frozen_error # also override dict methods __setitem__ = __delitem__ = clear = pop = popitem = setdefault = update = _frozen_error # alias mset = multiset # "mset" is a shortcut for "multiset" 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): if not isinstance(s, multiset): s = multiset.from_seq(s) if k is None: k = len(s) return mP(s, k) # a simple implementation of derangements (of distinct elements) # as in itertools elements are permuted by index, not value def derangements(s, k=None): s = list(s) if k is None: k = len(s) for p in itertools.permutations(range(len(s)), k): if not any(i == x for (i, x) in enumerate(p)): yield tuple(s[x] for x in p) # derangements where no value is equal to the value in the corresponding position of the original sequence def uD(s, k=None): s = list(s) for p in mpermutations(s, k): if not any(s[i] == x for (i, x) in enumerate(p)): yield p # cache is enabled by default @static(cache_enabled=1, cache={0: 1}) def subfactorial(n, validate=0): """subfactorial(n) is the number of derangements of n distinct elements""" if validate: n = as_int(n, include="0+") r = subfactorial.cache.get(n) if r is None: r = n * subfactorial(n - 1) + (-1 if n % 2 else 1) if subfactorial.cache_enabled: subfactorial.cache[n] = r return r # 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) sequence. '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 'uD' = unique derangements '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): try: select = subsets.select_fn[select] except KeyError: # [Python 3] raise ValueError(...) from None x = ValueError(str.format("invalid parameter: select={select!r}", select=select)) if hasattr(x, '__cause__'): x.__cause__ = None raise x # 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(s if isinstance(s, multiset) else multiset.from_seq(s)))), ('mP', mP, (lambda s: (s if isinstance(s, multiset) else multiset.from_seq(s)))), ('uD', uD, None), ): 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 def anagram(s): "generate anagrams of (treated as a string)" return subsets(str(s), size=len, select='mP', fn=join) # 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) @static(rtype=None) def filter2(p, seq, fn=list): """ use a predicate to partition a sequence 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 seq) if filter2.rtype is None: filter2.rtype = namedtuple('Filter2', 'true false') return filter2.rtype(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) @static(rtype=None) 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() >>> vs = [(1, 1), (1, 3), (2, 2), (3, 1), (3, 2), (3, 3)] "If I told you the first number you could deduce the second" >>> filter_unique(vs, (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(vs, (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) #printf("unq = {unq}\nnon = {non}") if filter_unique.rtype is None: filter_unique.rtype = namedtuple('FilterUnique', 'unique non_unique') return filter_unique.rtype(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(str.format("collect: failed to collect item: {x}", x=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 gfind(g, v, contains=operator.contains, fn=None): """ find the keys in group that contain the value . should map keys to a collection of values , and the default test is [[ v in vs ]], although this can be changed by specifying the parameter. by default a generator is returned, but this can be modified by the parameter. for instance to find the first value: [[ gfind(g, v, fn=peek) ]] or [[ peek(gfind(g, v)) ]] for all values as a set: [[ gfind(g, v, fn=set) ]] or [[ set(gfind(g, v)) ]] the function is also provided as: group.find(). """ ks = (k for (k, vs) in g.items() if contains(vs, v)) return (ks if fn is None else fn(ks)) @static(find=gfind) 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 # by considering values of sequence , we consider it to consist # of groups formed by contiguous subsequences with the same value of # we then look for the first value f(x) in the sequence that satisfies # function , and return (in order) all values in that group that # satisfy . def first_group(seq, fn=true, by=item(0), f=item(1)): (seq, g) = (iter(seq), None) # find the first value to satisfy fn for x in seq: v = f(x) if fn(v): yield v g = by(x) break else: # loop not terminated by break, so was not satisfied return # now look for remaining values in the same group for x in seq: if by(x) != g: return v = f(x) if fn(v): yield v # 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) unzip.__doc__ = "the inverse of zip" # cartesian product of a sequence, cproduct = unpack(itertools.product) def cproduct(ss, **kw): """ the cartesian product of a collection of collections. 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 """ fn = kw.pop('fn', None) rs = itertools.product(*ss, **kw) return ((fn(x) for x in rs) if fn else rs) # 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, env=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 = str.format("lambda {args}: {expr}", args=args, expr=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 = str.format("lambda *_x_: next({expr} for [{args}] in [_x_])", expr=expr, args=args) if env is None: # evaluate in the parent scope frame = sys._getframe(1) env = (frame.f_globals, frame.f_locals) return eval(expr, env[0], env[1]) # count the number of occurrences of a predicate in an iterator def icount(seq, p=None, t=None): """ count the number of elements in that satisfy predicate

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

use: icount(seq, p, n + 1) == n which is what icount_exactly(seq, 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(seq, p, n) == n This is what icount_at_least(seq, 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

use: icount(seq, p, n + 1) < n + 1 This is what icount_at_most(seq, p, n) does. The following will stop testing at 73 (the 21st prime): >>> icount_at_most(irange(1, 100), is_prime, 20) False If

and value , and it returns a function roots(n, k) that calculates the roots mod ^ e.g. to find the square roots of 9 mod 5488: >>> roots = poly_roots_mod(sq, 9) >>> sorted(roots(5488)) [3, 683, 2061, 2741, 2747, 3427, 4805, 5485] """ f = (p if v == 0 else (lambda x: p(x) - v)) cache = (dict() if cache else None) # find roots of f mod (b^k) using Hensel's lemma def hensel(b, k): if k == 0: return [0] assert k > 0 bk = b**k # check the cache if cache: rs = cache.get(bk) if rs is not None: return rs # find the roots k1 = k - 1 bk1 = b**k1 rs = list() for r in hensel(b, k1): for i in irange(b): x = i * bk1 + r if f(x) % bk == 0: rs.append(x) # cache if necessary if cache is not None: cache[bk] = rs return rs # find roots of f mod (n^k) def roots(n, k=1, fn=prime_factor): # check the cache if cache is not None: nk = (n if k == 1 else n**k) rs = cache.get(nk) if rs is not None: return rs # find roots for each prime factor of n (hs, bs) = (list(), list()) for (f, e) in fn(n): ek = (e if k == 1 else e * k) hs.append(hensel(f, ek)) bs.append(f**ek) # combine roots using CRT rs = list(crt(zip(xs, bs)).x for xs in cproduct(hs)) # cache if necessary if cache is not None: cache[nk] = rs return rs roots.p = p return roots # this may be faster than sqrtmod() for large (composite) m def poly_roots_mod_sqrtmod(a, m, k=1, fn=prime_factor): "find square roots of mod ^" return poly_roots_mod(sq, a)(m, k, fn) poly_roots_mod.sqrtmod = poly_roots_mod_sqrtmod # find square roots of mod ^ # this is OK for relatively small m, but more efficient (and complex) # approaches are available (e.g. pells.sqrtmod) def sqrtmod(a, m, k=1, lower=0): """ find square roots of mod . i.e. values such that ( * ) is congurent to (mod ). if 'lower' is set then only values <= m/2 are returned. see also: pells.sqrtmod() (or pells.sqrtmodp()) for a more efficient implementation. >>> sorted(sqrtmod(1, 16)) [1, 7, 9, 15] >>> sorted(sqrtmod(17, 43)) [19, 24] >>> sorted(sqrtmod(-1, 25)) [7, 18] """ if k > 1: m **= k a %= m if lower: for x in irange(0, m // 2): if (x * x) % m == a: # x is a root yield x else: for x in irange(0, m // 2): if (x * x) % m == a: # x is a root yield x # -x (mod m^k) is also a root if x > 0 and m > 2 * x: yield m - x # solve linear diophantine equations in 2 variables: a.X + b.Y = v def diop_linear(a, b, v=0, mX=0, fn=0): """ solve the linear Diophantine equation a.X + b.Y = k for integers X, Y (where a, b are non-zero, and gcd(a, b) divides k). return ((X0, Y0), (Xd, Yd)) to give solutions of the form: (X, Y) = (X0 + t.Xd, Y0 + t.Yd) for integer t the value of X0 returned is the smallest integer possible above (or equal to) mX, and Xd is positive. however, if is set, then a function f: t -> (X, Y) is returned instead. """ if a == 0 or b == 0: raise ValueError("diop_linear: invalid equation") (X, Y, g) = egcd(a, b) if g > 1: (a, b, v) = (a // g, b // g, div(v, g)) if v is None: raise ValueError("diop_linear: no solutions") # calculate particular solution (X0, Y0) and deltas (Xd, Yd) (X0, Y0) = (v * X, v * Y) (Xd, Yd) = ((-b, a) if b < 0 else (b, -a)) # adjust X0 to be the smallest value t = divc(mX - X0, Xd) X0 += t * Xd Y0 += t * Yd #assert all(a * (X0 + t * Xd) + b * (Y0 + t * Yd) == c for t in irange(-50, 50)) if fn: return (lambda t, X0=X0, Y0=Y0, Xd=Xd, Yd=Yd : (X0 + t * Xd, Y0 + t * Yd)) return ((X0, Y0), (Xd, Yd)) # solve linear diophantine equation in 3 variables: a.X + b.Y + c.Z = v def diop_linear3(a, b, c, v=0, rX=(0, inf), rY=(0, inf), rZ=(0, inf)): ((min_X, max_X), (min_Y, max_Y), (min_Z, max_Z)) = (rX, rY, rZ) g = gcd(a, b) (a_, b_) = (a // g, b // g) # solve: g.U + c.Z = v min_U = a_ * (min_X if a_ > 0 else max_X) + b_ * (min_Y if b_ > 0 else max_Y) max_U = a_ * (max_X if a_ > 0 else min_X) + b_ * (max_Y if b_ > 0 else min_Y) f1 = diop_linear(g, c, v, mX=min_U, fn=1) for i in irange(0, inf): (U, Z) = f1(i) if U > max_U: break if Z < min_Z or Z > max_Z: continue # solve: a_.X + b_.Y = U f2 = diop_linear(a_, b_, U, mX=min_X, fn=1) for j in irange(0, inf): (X, Y) = f2(j) if X > max_X: break if Y < min_Y or Y > max_Y: continue yield (X, Y, Z) # 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 """ rest = set(rest) rest.discard(a) return reduce(gcd, rest, a) mgcd = (getattr(math, 'gcd', _mgcd) if _pythonv > (3, 8) else _mgcd) # 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 """ rest = set(rest) rest.discard(a) return reduce(lcm, rest, a) mlcm = (getattr(math, 'lcm', _mlcm) if _pythonv > (3, 8) else _mlcm) def mlcmq(*qs): """ LCM of multiple fractions (each represented as (, )) >>> mlcmq((1, 2), (2, 3)) (2, 1) """ (ns, ds) = zip(*qs) return fraction(mlcm(*ns), mgcd(*ds)) def is_coprime(*vs): """ check values are _pairwise_ coprime. if you want to check for _setwise_ coprime use: mgcd(*vs) == 1 >>> is_coprime(5, 6, 7) True >>> is_coprime(5, 6, 7, 8) False >>> mgcd(5, 6, 7, 8) 1 """ n = len(vs) if n < 2: return True if n == 2: return gcd(*vs) == 1 # equivalent to [[ return mlcm(*vs) == multiply(vs) ]] # calculate lcm and product m = p = 1 for v in vs: m = lcm(m, v) p *= v if m != p: return False return True # multiple divmod def mdivmod(x, *vs): """ perform multiple applications of divmod. examples: 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) >>> mdivmod(45296, 60, 60) [12, 34, 56] """ rs = list() for v in reversed(vs): (x, r) = divmod(x, v) rs.insert(0, r) rs.insert(0, x) return rs def qsum(fs, normal=1): """ determine the numerator and denominator of the sum of the fractions given in , a sequence of (, ) pairs. if normal is set, then the fraction will be reduced to lowest terms. fraction([(a, b), (c, d), (e, f), ...]) -> a/b + c/d + e/f + ... >>> qsum([(1, 2), (1, 3), (1, 6)]) # 1/2 + 1/3 + 1/6 = 1 (1, 1) >>> qsum([(1, 2), (3, 4), (5, 6)]) # 1/2 + 3/4 + 5/6 = 25/12 (25, 12) """ if not fs: return (0, 1) # = 0 # go through the fractions in pairs fs = iter(fs) (a, b) = next(fs) for (c, d) in fs: (a, b) = (a * d + b * c, b * d) if b == 0: raise ZeroDivisionError("fraction can't have zero denominator") if normal: if b < 0: (a, b) = (-a, -b) g = gcd(a, b) if g != 1: (a, b) = (a // g, b // g) return (a, b) # sum of reciprocals (= unit fractions) def rsum(ds, normal=1): """determine the sum of the reciprocals 1/d1 + 1/d2 + ...""" if not ds: return (0, 1) # = 0 ds = list(ds) if len(ds) == 1: return (1, ds[0]) b = multiply(ds) if b == 0: raise ZeroDivisionError("fraction can't have zero denominator") a = sum(b // d for d in ds) if normal: if b < 0: (a, b) = (-a, -b) g = gcd(a, b) if g != 1: (a, b) = (a // g, b // g) return (a, b) # for those times when Rational() is overkill def fraction(*args, **kw): """ 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 not args: return (0, 1) # = 0 return qsum(chunk(args, 2), normal=kw.get('normal', 1)) @static(rtype=None) def Fraction(*args): """ same as fraction(), but returns an object where the numerator and denominator can be referred to as obj.num and obj.den. """ (a, b) = fraction(*args) if Fraction.rtype is None: Fraction.rtype = namedtuple('Fraction', 'num den') return Fraction.rtype(a, b) def format_fraction(n, d, base=0): """format fraction as a string""" s = int2base(n, base=base) if d == 1: return s return s + "/" + int2base(d, base=base) def str2fraction(s, base=0): """parse a fraction from a string""" if base == 0: base = radix (ns, _, ds) = map(str.strip, str.partition(s, "/")) n = base2int(ns, base=base) if not ds: return (n, 1) d = base2int(ds, base=base) return fraction(n, d) 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)) def ratio_q(*qs): """ return ratio of fractions in as integers in lowest terms. >>> ratio_q(rational(2, 3), rational(10, 3)) (1, 5) """ # turn the fractions into integers m = call(mlcm, (q.denominator for q in qs)) return call(ratio, (int(m * q) for q in qs)) # check values in r1 are the same ratio as values in r2 def ratio_eq(r1, r2, **kw): """ >>> ratio_eq((3, 4), (6, 8)) True >>> ratio_eq((3, 4, 5), (6, 8, 10)) True """ strict = kw.pop('strict', None) rs = iter(zip(r1, r2, strict=strict) if strict is not None else zip(r1, r2)) # get the first pair try: (p, q) = next(rs) except StopIteration: return True # and check the rest # was: [[ return all(p * y == x * q for (x, y) in rs) ]] for (x, y) in rs: if p * y != x * q: return False return True fraction_eq = ratio_eq # 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 class Rdiv(object): """ create a function that will calculate a/b, and return an integer result if the result is an integer, or a rational object if the result is a rational """ 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() # harmonic sum def hsum(vs, div=rdiv, num=1): """ compute the harmonic sum of the specified values: e.g. hsum([a, b, c]) = 1/(1/a + 1/b + 1/c) """ return div(num, sum(div(1, v) for v in vs)) # arithmetic mean def amean(vs, div=rdiv): """ calculate the arithmetic mean of values . for integer values the result is generally rational, if you only want integer values use div=div. """ return avg(vs, div=div) # geometric mean def gmean(vs, root=root): """ calculate the geometric mean of values . for integer values the result is generally irrational (float), if you only want integer values use root=is_power (or: is_square for 2 arguments, is_cube for 3 arguments). """ p = 1 k = 0 for v in vs: p *= v k += 1 return (p if k == 1 else root(p, k)) # quadratic mean (aka root mean square) def qmean(vs, root=math.sqrt): return root(amean(sq(v) for v in vs)) # harmonic mean def hmean(vs, div=rdiv): """ calculate the harmonic mean of values . for integer values the result is generally rational, if you only want integer results use as_int(hmean(...)). """ return hsum(vs, div=div, num=len(vs)) @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, **kw): """ return a! (or a! / b! if b is specified). >>> factorial(6) 720 >>> factorial(10, 7) 720 """ validate = kw.pop('validate', 0) if kw: raise TypeError(str.format("factorial: unknown arguments {kw}", kw=seq2str(kw.keys()))) if validate: a = as_int(a, include="0+") if not bs: return math.factorial(a) if validate: bs = tuple(as_int(b, include="0+") for b in bs) 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 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 return math.factorial(n) // math.factorial(r) // math.factorial(n - r) C = nCr # also known as: "stars and bars" (or "sheep and fences") # 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 "If each bag contains 122 sweets, and they come in 5 different flavours how many different packets can there be?" >>> int2words(M(5, 122)) 'ten million nine thousand one hundred and twenty five' """ return C(n + k - 1, k) def recurring_digits(a, b, recur=0, base=0): """ find the non-recurring/recurring digits in the base expansion of the proper fraction / . """ if base == 0: base = radix 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))) r = dict() s = list() n = 0 while True: j = r.get(a) # have we seen this dividend before? if j is not None: return (s[:j], s[j:]) # 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): s.append(d) @static(rtype=None) def recurring(a, b, recur=0, base=0, digits=None): """ find recurring representation of the fraction / in the specified base. returns strings (, , ) if you want rationals that normally terminate represented as non-terminating set to a true value. >>> tuple(recurring(1, 7)) ('0', '', '142857') >>> tuple(recurring(3, 2)) ('1', '5', '') >>> tuple(recurring(593, 530)) ('1', '1', '1886792452830') >>> tuple(recurring(3, 2, recur=1)) ('1', '4', '9') >>> tuple(recurring(5, 17, base=16)) ('0', '', '4B') """ if base == 0: base = radix # 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) i2b = partial(int2base, base=base, digits=digits) i = i2b(i) (nr, rr) = (str.join('', map(i2b, ds)) for ds in recurring_digits(a, b, recur=recur, base=base)) if neg and (nr or rr or i != '0'): i = '-' + i if recurring.rtype is None: recurring.rtype = namedtuple('Recurring', 'i nr rr') return recurring.rtype(i, nr, rr) # 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.pop('dp', '.') enc = kw.pop('enc', '()') dots = kw.pop('dots', '...') if kw: raise TypeError(str.format("format_recurring: unknown arguments {kw}", kw=seq2str(kw.keys()))) 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=0, 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) """ if base == 0: base = radix (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)) # thanks to Brian for testing this routine def reciprocals(k, b=1, a=1, m=1, M=inf, g=0, rs=[]): """ see also: sum_of_reciprocals(). generate positive integers (d1, d2, ..., dk) such that 1/d1 + 1/d2 + ... + 1/dk = a/b. the denominators are generated as an ordered list m = minimum allowed denominator M = maximum allowed denominator g = minimum allowed gap between denominators 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]] >>> list(reciprocals(3, 3, M=15)) [[5, 15, 15], [6, 10, 15], [6, 12, 12], [8, 8, 12], [9, 9, 9]] there are 3462 ways to express 1 as the sum of 6 reciprocals [OEIS A002966] >>> icount(reciprocals(6)) 3462 or 2320 if the reciprocals must all be different [OEIS AA006585] >>> icount(reciprocals(6, g=1)) 2320 """ if a < 1 or b < 1 or k < 1 or m < 1 or M < m or g < 0: return # check remaining fraction against the k largest possible reciprocals if g == 0 or k < 2: z = m * a - k * b if z == 0: yield rs + ([m] * k) else: x = m + (k - 1) * g if x > M: return ds = list(irange(m, x, step=g)) (p, q) = rsum(ds, normal=0) z = q * a - p * b if z == 0: yield rs + ds if z >= 0: return # more to do? if k > 2: dmin = divc(b + 1, a) dmax = divf(k * b, a) if M != inf: # if M is given we can maybe find a better dmin [suggested by Frits] x = M - g * (k - 2) if x < m: return (p, q) = ((k - 1, M) if g == 0 else rsum(irange(x, M, step=g), normal=0)) z = q * a - p * b if z > 0: dmin = max(dmin, divc(b * q, z)) # find a suitable next reciprocal for d in irange(max(m, dmin), min(M, dmax)): # solve for the remaining fraction #yield from reciprocals(k - 1, b * d, a * d - b, d + g, M, g, rs + [d]) #[Python 3] for ds in reciprocals(k - 1, b * d, a * d - b, d + g, M, g, rs + [d]): yield ds #[Python 2] elif k == 2: # special case k = 2 dmin = divc(b + 1, a) dmax = divf(b + b, a) if M != inf: x = a * M - b if x < 1: return dmin = max(dmin, divc(b * M, x)) 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] elif k == 1: (d, r) = divmod(b, a) if r == 0 and not (d < m or d > M): yield rs + [d] # alternative interface to reciprocals # allows f can be specified as (a, b) for a/b or just b for 1/b # allows input validation def sum_of_reciprocals(k, f, min_d=1, max_d=inf, gap_d=0, validate=0): """ generate k-sets of reciprocals (unit fractions) that sum to the fraction f. f can be specified as a single positive integer to find solutions for the unit fraction 1/f, or as a pair of positive integers (a, b) to find solutions for the fraction a/b. return values are collections of k positive integers (d1, d2, ..., dk), returned as an ordered list, such that: f = 1/d1 + 1/d2 + ... + 1/dk min_d = minimum allowed denominator max_d = maxumum allowed denominator gap_d = minimum allowed gap between denominators gap_d=1 gives distinct denominators; gap_d=0 allows repeated values validate = enable parameter validation """ try: (a, b) = f except TypeError: (a, b) = (1, f) if validate: (k, b, a, m) = (as_int(x, include="+") for x in (k, b, a, m)) if M != inf: M = as_int(M, include="+") g = as_int(g, include="0+") return reciprocals(k, b, a, m=min_d, M=max_d, g=gap_d) # 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: # TODO: decode command line args to unicode if args is None: args = sys.argv[1:] get_argv.argv = args 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 (if vs is a dict) def __sprintf(fmt, vs): return str.format(fmt, **vs) # Python 3 has str.format_map(vs) def __sprintf3(fmt, vs): return str.format_map(fmt, 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), None, vs) class dictify(object): def __init__(self, fn): self.fn = fn def __getitem__(self, key): return self.fn(key) @static(fn=None) def _sprintf(fmt, fn=None, vs=None, frame=None): dstack = getattr(collections, 'ChainMap', None) if dstack: # in Python 3 we can use collections.ChainMap ds = [] if fn: ds.append(dictify(fn)) if vs: ds.append(vs) if frame: ds.append(frame.f_locals) ds.append(frame.f_globals) d = (ds[0] if len(ds) == 1 else dstack(*ds)) else: # we can populate a dict d = dict() if frame: d.update(frame.f_globals) d.update(frame.f_locals) if vs: d.update(vs) if fn: print("WARNING: sprintf 'fn' parameter is not supported in Python {_pythonv}".format(_pythonv=join(_pythonv, sep='.'))) return _sprintf.fn(fmt, d) if _pythonv > (3, 5): _sprintf.fn = __sprintf36 elif _python > 2: _sprintf.fn = __sprintf3 else: _sprintf.fn = __sprintf def strfmt(fmt='', fn=None, vs=None, frame=None): """ format the string using substitutions from: - function (in Python 3) for dynamic lookup - dictionary - local/global variables in frame """ if isinstance(frame, int): frame = sys._getframe(frame) return _sprintf(fmt, fn, vs, frame) # print with variables interpolated into the format string def sprintf(fmt='', **vs): """ 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, None, vs, sys._getframe(1)) # print with local variables interpolated into the format string def printf(fmt='', **vs): """ print format string with interpolated local variables and keyword arguments. the final newline can be suppressed by ending the string with '\\' (which you may need to escape). >>> (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, None, vs, sys._getframe(1)) d = dict() # flush=1 if s.endswith('\\'): (s, d['end']) = (s[:-1], '') print(s, **d) # format output into blocks class Output(): def __init__(self, start=None, end="", prefix=" ", timed=None, override=1): self.start = start self.end = end self.prefix = prefix self.timer = timed self.override = override self.held = list() self.immediate = 0 self._saved = None def __enter__(self): start = self.start if start is not None: printf("{start}") if self.override: # override names: printf -> self.printf frame = sys._getframe(1) if frame: vs = frame.f_locals d = dict() for k in ['printf']: fn = vs.get(k) if fn: d[k] = fn vs[k] = self.printf if d: self._saved = (vs, d) if self.timer: kw = dict() if isinstance(self.timer, basestring): kw['timer'] = self.timer self.timer = Timer(**kw) return self def __exit__(self, *args): timer = self.timer if timer: timer.stop() timer.report() if self._saved: # restore any saved names (vs, d) = self._saved for (k, v) in d.items(): vs[k] = v self._saved = None end = self.end if end is not None: printf("{end}") def printf(self, fmt='', **kw): printf("{prefix}{s}", prefix=self.prefix, s=_sprintf(fmt, None, kw, sys._getframe(1))) def holdf(self, fmt='', **kw): "this allows you to delay the printing of a string until release() is called" self.held.append(_sprintf(fmt, None, kw, sys._getframe(1))) if self.immediate: self.release() def release(self): "output strings previously delayed using holdf()" held = self.held if held: for s in held: printf("{prefix}{s}", prefix=self.prefix) self.held = list() def discard(self): "discard any previously delayed strings" self.held = list() @classmethod def partial(cls, **kw): return functools.partial(cls, **kw) # to use: output.printf() (and more probably output.holdf(), output.release()) output = Output(prefix='') class Failure(Exception): pass # fatal error @static(Failure=Failure) def fail(expr=True, msg=''): """ if is true then a Failure exception is raised (with optional message ). (NOTE: "assert" statements raise an exception when the expression is false). otherwise the value of is returned (which will be a false value), this allows constructions of the form: fail(expr1) or expr2 which will fail if is true, but if it is false it will return the value of or: expr2 if expr1 else fail() which will fail if is false, but if it is true it will return the value of """ if expr: raise Failure(msg) return expr # non-fatal warning def warn(expr=True, msg=''): if expr: print("WARNING:", msg) return expr def catch(*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 """ (fn, args) = (args[0], args[1:]) try: return fn(*args, **kw) except Exception: if 'v' in _PY_ENIGMA: print("catch: caught exception!") print(sys.exc_info()) return # round endpoints to integers def irange_round(a, b, step=1, rnd='i'): """ round the endpoints of a range to integer values. the type of rounding used can be specified using the 'rnd' parameter: i = round endpoints to give the largest range that is inside the given range (default) x = round endpoints to give the smallest range that includes the given range I = round endpoints to give the largest range of multiples of inside the given range X = round endpoints to give the smallest range of multiples of that includes the given range f = use floor rounding (intf) c = use ceiling rounding (intc) r = use nearest integer rounding (intr) b = use Python builtin rounding (round) F = use floor rounding to the nearest multiple of C = use ceiling round to the nearest multiple of """ # f = floor, c = ceil, r = round m = abs(step) fn = dict(f=intf, c=intc, r=intr, b=round, F=partial(floor, m=m), C=partial(ceil, m=m), U=identity) (ka, kb) = (rnd[0], rnd[-1]) # i = internal, x = external if ka == 'i': ka = ('c' if step > 0 else 'f') elif ka == 'x': ka = ('f' if step > 0 else 'c') elif ka == 'I': ka = ('C' if step > 0 else 'F') elif ka == 'X': ka = ('F' if step > 0 else 'C') if b == inf or b == -inf: kb = 'U' elif kb == 'i': kb = ('f' if step > 0 else 'c') elif kb == 'x': kb = ('c' if step > 0 else 'f') elif kb == 'I': kb = ('F' if step > 0 else 'C') elif kb == 'X': kb = ('C' if step > 0 else 'F') return irange(fn[ka](a), fn[kb](b), step=step) # irange_ndigits(4) = (1000, 1001, ..., 9999) # irange_ndigits(2, 3) = (10, 11, ..., 999) def irange_ndigits(a, b=None, step=1, base=0): if base == 0: base = radix if b is None: b = a # smallest number with digits, largest number with digits (a, b) = ((0 if a == 1 else base**(a - 1)), (base**b) - 1) if step < 0: (a, b) = (b, a) return irange(a, b, step=step) # find the common difference of an arithmetic sequence (or None) def aseq_diff(seq): return seq_all_same_r(y - x for (x, y) in tuples(seq, 2)).value # calculate the sum of an arithmetic sequence # (for large ranges this is faster than calling sum() on it) def aseq_sum(seq): n = catch(len, seq) if n is None: return None return (0 if n == 0 else n * (seq[0] + seq[-1]) // 2) # inclusive range iterator # irange(a, b) -> [a, a + 1, ..., b] # irange(n) -> irange(0, n - 1) -> [0, ..., n - 1] @static(inf=inf, round=irange_round, ndigits=irange_ndigits, sum=(lambda *args, **kw: aseq_sum(irange(*args, **kw)))) # so irange. can be used def irange(a, b=None, step=1): """ irange(a, b) = an integer range iterator that includes both endpoints, and . it will generate, in order, the integers: [a, a + k, a + 2k, ..., b] where is the step. note that it is possible to choose endpoint/step combinations where the sequence of integers generated does not include , or is empty. if is specified as inf (or -inf for negative steps) the iterator will generate values indefinitely. if you want the endpoints to be rounded to integers use irange.round(). 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 not step: raise ValueError("irange: step cannot be zero") if a == inf or a == -inf: raise ValueError("irange: unbounded start value") 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 b - 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 # conditional range def crange(a=0, fn=true, step=1): """ generate the sequence: [a, a + step, a + 2*step, ...] until the 'fn' condition returns false. the value that fails the 'fn' condition is not returned. >>> list(crange(0, lt(10))) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] >>> list(crange(0, le(10))) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] >>> list(crange(10, ge(0), step=-1)) [10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0] """ x = a while fn(x): yield x x += step # flatten a list of lists # [Python 3.15: has unpacking in comprehensions] def flatten(seq, skip=1, fn=list): """ flatten a sequence of sequences. if 'skip' is set then iterators that evaluate to 'False' are ignored (i.e. treated as the empty sequence). 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'] """ if skip: seq = filter(None, seq) r = (x for itr in seq for x in itr) # = (*itr for itr in seq) for Python 3.15+ return (r if fn is None else fn(r)) # similar to using: #return fn(itertools.chain.from_iterable(seq)) # NOTE: but the following is not necessarily equivalent if generators are involved #return fn(itertools.chain(*seq)) # shortcut for flatten() that returns a generator iflatten = partial(flatten, fn=None) # 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) """ if 'fn' not in kw: kw['fn'] = None return flatten(ss, **kw) # generate permutations of the items of a sequence def permute(ss, size=len, select='P'): for s in ss: #yield from subsets(s, size=size, select=select) #[Python 3] for z in subsets(s, size=size, select=select): yield z #[Python 2] # flatten(zip(*ss), fn=iter) works if arguments are the same length def interleave(*ss): """interleave values from a collection of iterators""" 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 - a sequence 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 return _flattened(z, depth, test) # merge a collection of sorted iterables into a single iterable def merge(seqs, uniq=0): from heapq import merge r = merge(*seqs) return (uniq1(r) if uniq else r) # 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. if the container has a sense of order, items are added at the end. >>> 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): r = type(s)(s) r.extend(vs) return r if isinstance(s, tuple): return s + vs if isinstance(s, basestring): return s + str.join('', vs) if isinstance(s, set): r = s.copy() r.update(vs) return r if isinstance(s, fset): return s.union(vs) if isinstance(s, multiset): return s.copy().update_from_seq(vs) raise ValueError(str.format("append() can't handle container of type {x}", x=type(s))) # this unifies removing an element (or elements) from a container def remove(s, *vs, **kw): """ make a new container, the same as but with values in removed. items appearing in that are not in are ignored. if the container has a sense of order, the earliest items are removed >>> remove((1, 2, 3), 2) (1, 3) >>> remove([1, 2, 3, 2], 2) [1, 3, 2] >>> remove({1, 2, 3}, 2) == {1, 3} True >>> remove('1232', '2') '132' """ strict = kw.pop('strict', 0) if kw: raise TypeError(str.format("remove: unknown arguments {kw}", kw=seq2str(kw.keys()))) if isinstance(s, list): r = type(s)(s) for v in vs: try: r.remove(v) except ValueError as e: if strict: raise ValueError(str.format("{v} not in list", v=v)) continue return r if isinstance(s, tuple): return type(s)(remove(list(s), *vs)) if isinstance(s, basestring): for v in vs: if strict and v not in s: raise ValueError(str.format("{v!r} not in string", v=v)) s = s.replace(v, '', 1) return s if isinstance(s, (set, fset, multiset)): if strict and not s.issuperset(vs): raise ValueError(str.format("elements not in *set")) return s.difference(vs) raise ValueError(str.format("remove() can't handle container of type {x}", x=type(s))) # remove all (possibly none) values in from a container def discard(s, *vs): if isinstance(s, (list, tuple)): fn = type(s) if len(vs) == 1: v = vs[0] return fn(x for x in s if x != v) else: vs = set(vs) return fn(x for x in s if x not in vs) if isinstance(s, (set, multiset, fmultiset)): r = s.copy() for v in vs: r.discard(v) return r if isinstance(s, basestring): r = s for v in vs: r = r.replace(v, '') return r raise ValueError(str.format("discard() can't handle container of type {x}", x=type(s))) # swap items in container at keys and def swap(s, i, j, *ks, **kw): """ create a container the same as but with keys and swapped. """ fn = kw.get('fn', None) if isinstance(s, Mapping): r = type(s)(s) elif isinstance(s, (list, tuple, basestring)): r = list(s) if fn is None: fn = (''.join if isinstance(s, basestring) else type(s)) (r[i], r[j]) = (r[j], r[i]) if ks: for (i, j) in chunk(ks, 2): (r[i], r[j]) = (r[j], r[i]) return (r if fn is None else fn(r)) # restriction of a container that only includes items at keys def restrict(s, ks, strict=0): """ create the restriction of container that only includes keys in keys that do not occur in the original container are ignored, unless the flag is set, in which case an exception is raised. >>> map2str(restrict(dict(a=1, b=2, c=3, d=4), {'a', 'b', 'x'})) '(a=1, b=2)' >>> restrict(['zero', 'one', 'two', 'three'], [1, 3]) ['one', 'three'] >>> restrict("abracadabra", [0, 3, 5, 7, 10]) 'aaaaa' """ if isinstance(s, Mapping): r = type(s)() for k in ks: try: r[k] = s[k] except KeyError: if strict: raise return r if isinstance(s, (list, tuple, basestring)): r = list() for k in ks: try: r.append(s[k]) except IndexError: if strict: raise return (str.join('', r) if isinstance(s, basestring) else type(s)(r)) # 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)] """ if n < 1: raise ValueError(str.format("invalid tuple length: {n!r}", n=n)) seq = iter(seq) if circular and n > 1: # we need extract the first (n - 1) items and add them to the end xs = first(seq, n - 1) if not xs: return m = len(xs) if m < n - 1: seq = itertools.chain(xs, seq, (xs[k % m] for k in xrange(n - 1))) else: seq = itertools.chain(xs, seq, xs) t = list() try: # collect the first tuple for _ in xrange(n): t.append(next(seq)) while True: # return the tuple yield fn(t) # move the next value in to the tuple t.pop(0) t.append(next(seq)) except StopIteration: return def last(seq, count=1, fn=list): """ find the last items in sequence . >>> last([1, 2, 3, 4]) [4] >>> last(primes.between(1, 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 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 call(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, max_q=inf, validate=0): """ express total using denominations . optional: using minimum quantity (integer), and maximum quantity (integer, or inf); or: quantities chosen from (if is specified the and arguments are ignored). and should be increasing sequences. (this can be checked with 'validate=1'). generated values are the quantities for each denomination in . >>> sorted(express(20, [3, 5, 7])) [(0, 4, 0), (1, 2, 1), (2, 0, 2), (5, 1, 0)] >>> sorted(express(20, [3, 5, 7], min_q=1)) [(1, 2, 1)] >>> sorted(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) if validate: # check ds (and qs) for strictly increasing sequences if not is_increasing(ds, strict=1): raise ValueError(str.format("express: invalid denominations {ds!r}", ds=ds)) if qs is not None and ((not qs) or (not is_increasing(qs, strict=1))): raise ValueError(str.format("express: invalid quantities {qs!r}", qs=qs)) if not (ds and ds[0] > 0): raise ValueError(str.format("invalid denominations {ds!r}", ds=ds)) if qs: return express_quantities(t, ds, qs) if min_q != 0: return express_denominations_min(t, ds, min_q, max_q) return express_denominations(t, ds, max_q=max_q) # express total using denominations with results in # working backwards from index def _express_denominations(t, ds, ss, i, max_q=inf): if t == 0: if i >= 0: ss[:i + 1] = [0] * (i + 1) yield ss else: d = ds[i] (k, r) = divmod(t, d) if i == 0: if r > 0 or k > max_q: return ss[0] = k yield ss else: for q in irange(0, min(k, max_q)): ss[i] = q #yield from _express_denominations(t - d * q, ds, ss, i - 1, max_q) #[Python 3] for z in _express_denominations(t - d * q, ds, ss, i - 1, max_q): yield z #[Python 2] # express total using denominations def express_denominations(t, ds, max_q=inf): n = len(ds) for ss in _express_denominations(t, ds, [0] * n, n - 1, max_q): yield tuple(ss) # express total using denominations , min quantity def express_denominations_min(t, ds, min_q, max_q): # allocate the minimum quantities t -= min_q * sum(ds) n = len(ds) if t == 0: yield (min_q,) * n elif t > 0: # solve for the remaining amount for ss in _express_denominations(t, ds, [0] * n, n - 1, max_q=max_q - min_q): # add in the initial quantities yield tuple(q + min_q for q in ss) # express total using denominations , quantities chosen from def _express_quantities(t, ds, qs, ss, i): if t == 0: if i >= 0: if 0 not in qs: return ss[:i + 1] = [0] * (i + 1) yield ss else: (d, qs_) = (ds[i], qs) if i == 0: (k, r) = divmod(t, d) if r or k not in qs: return qs_ = [k] for q in qs_: t_ = t - d * q if t_ < 0: break ss[i] = q #yield from _express_quantities(t_, ds, qs, ss, i - 1) #[Python 3] for z in _express_quantities(t_, ds, qs, ss, i - 1): yield z #[Python 2] # deal with non-zero min quantity def express_quantities(t, ds, qs): (min_q, n) = (qs[0], len(ds)) if min_q == 0: for ss in _express_quantities(t, ds, qs, [0] * n, n - 1): yield tuple(ss) else: # allocate the minimum quantities t -= min_q * sum(ds) if t == 0: yield (min_q,) * n elif t > 0: # solve for the remaining amount for ss in _express_quantities(t, ds, tuple(q - min_q for q in qs), [0] * n, n - 1): yield tuple(q + min_q for q in ss) # express using (, ) pairs (see: multiset.express()) def _express_pairs(t, vs, tv, ss, i): if t == 0: if i >= 0: ss[:i + 1] = [0] * (i + 1) yield ss elif t > tv or i < 0: return # if i == 1, we could use diop_linear() to determine n0 and n1 (but it is generally slower) elif i == 0: # 1 denomination remaining (x, q) = vs[0] (n, r) = divmod(t, x) if not (r != 0 or n > q): ss[0] = n yield ss else: (x, q) = vs[i] if x > t: if tv < inf: tv -= x * q ss[i] = 0 #yield from _express_pairs(t, vs, tv, ss, i - 1) # [Python 3] for z in _express_pairs(t, vs, tv, ss, i - 1): yield z # [Python 2] else: max_n = min(q, t // x) if tv < inf: tv -= x * q for n in irange(0, max_n): ss[i] = n #yield from _express_pairs(t - n * x, vs, tv, ss, i - 1) # [Python 3] for z in _express_pairs(t - n * x, vs, tv, ss, i - 1): yield z # [Python 2] # express using (, ) pairs in exactly coins def _express_pairs_k(t, vs, tv, k, ss, i, d0, q0): if t == 0 or k == 0: if t == 0 and k == 0: if i >= 0: ss[:i + 1] = [0] * (i + 1) yield ss elif t > tv or i < 0 or t < k * d0: return elif i < 2: if i == 1: # 2 denominations remaining: n0 * d0 + n1 * d1 = t; n0 + n1 = k (d1, q1) = vs[1] (n0, r) = divmod(k * d1 - t, d1 - d0) if r == 0 and not (n0 < 0 or n0 > k or n0 > q0): n1 = k - n0 if not (n1 > q1): ss[0] = n0 ss[1] = n1 yield ss else: # i == 0 # 1 denomination remaining if (not k > q0) and t == k * d0: ss[0] = k yield ss elif k == 1: # 1 coin remaining, just search for it for j in irange(i, 0, step=-1): (d, _) = vs[j] if d == t: ss[:i + 1] = [0] * (i + 1) ss[j] = 1 yield ss break if d < t: break else: (d, q) = vs[i] if k <= q and t == k * d: ss[i] = k if i > 1: ss[:i] = [0] * i yield ss elif d > t: if tv < inf: tv -= d * q ss[i] = 0 #yield from _express_pairs_k(t, vs, tv, k, ss, i - 1, d0, q0) # [Python 3] for z in _express_pairs_k(t, vs, tv, k, ss, i - 1, d0, q0): yield z # [Python 2] else: max_n = min(q, t // d) if k * d < t: return # target cannot be achieved in k coins [suggested by Frits] if k < max_n: max_n = k if tv < inf: tv -= d * q for n in irange(0, max_n): ss[i] = n #yield from _express_pairs_k(t - n * d, vs, tv, k - n, ss, i - 1, d0, q0) # [Python 3] for z in _express_pairs_k(t - n * d, vs, tv, k - n, ss, i - 1, d0, q0): yield z # [Python 2] def express_pairs(t, vs, tv=inf, k=None, validate=0): """ express total using (, ) pairs vs = ordered list of (, ) pairs tv = total sum of values in (= inf if any is inf) k = express using values (not _different_ denominations) validate = run sanity checks on parameters note that values and can be inf. returns a list of (, ) pairs """ if validate: # sanity check parameters vs = sorted((d, q) for (d, q) in vs if q > 0 and d <= t) if not is_increasing(vs, strict=1): raise ValueError("express_pairs: invalid parameter") if tv < sum(d * q for (d, q) in vs): raise ValueError("express_pairs: invalid parameter") if t < 0: raise ValueError("express_pairs: invalid parameter") if k is not None and k < 0: raise ValueError("express_pairs: invalid parameter") n = len(vs) if k is None: sss = _express_pairs(t, vs, tv, [0] * n, n - 1) else: (d0, q0) = vs[0] sss = _express_pairs_k(t, vs, tv, k, [0] * n, n - 1, d0, q0) for ss in sss: yield list((d, q) for ((d, _), q) in zip(vs, ss) if q > 0) # 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(str.format("invalid denominations: {denominations}", 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]) # 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(ns) 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))): #yield from _decompose(t - n, k_, M(n, d), max_v, d, R, M, r, fn, ns + (n,)) #[Python 3] for z in _decompose(t - n, k_, M(n, d), max_v, d, R, M, r, fn, ns + (n,)): yield z #[Python 2] # 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 = minimum separation between numbers in the sequence; 0 allows repeats [default: 1] min_v = minimum permissible value (non-negative integer) [default: 1] max_v = maximum permissible value (non-negative integer, or inf) [default: inf] fn = return type (default is to return tuples) Note that the 'sep' parameter specifies the separation between the numbers used in constructing the sequence, not the separation between adjacent numbers in the returned sequence. So when using 'increasing=0' specifying 'sep=1' means the numbers will all be different, and 'sep=0' will allow repeats. """ if increasing == 0: if sep == 0: # generate sequences, allowing repeats R = (lambda t, k, k_, m: t - k_ * m) M = (lambda n, d: min_v) return (lambda t, k=k, min_v=min_v, max_v=max_v, fn=fn: _decompose(t, k, min_v, max_v, 0, R, M, 0, fn)) else: # generate increasing sequences with the appropriate sep value and then permute the answers f = Decompose(k, increasing=1, sep=sep, min_v=min_v, max_v=max_v, fn=fn) return (lambda t, k=k: permute(f(t, k))) 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, max_v=max_v, fn=fn: _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 into -sequences of non-negative integers that sum to t = total sum of each sequence (or a sequence of totals, used in turn) 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) See also: Decompose(). >>> 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)] >>> sorted(decompose(6, 3, increasing=0, sep=1, min_v=1)) [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)] """ f = Decompose(increasing=increasing, sep=sep, min_v=min_v, max_v=max_v, fn=fn) # is the total iterable? try: ts = iter(t) except TypeError: ts = (t,) for t in ts: #yield from f(t, k) #[Python 3] for z in f(t, k): yield z #[Python 2] ############################################################################### # exact set cover (using Knuth's Algorithm X) # in-place algorithm X 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() # find exact hitting sets using algorithm X def exact_hitting_set(ss): """ find exact hitting sets for the sets in . each returned value consists of a collection of elements where every set from has exactly one element in common with the returned value. # from [ https://en.wikipedia.org/wiki/Exact_cover#Exact_hitting_set ] >>> list(exact_hitting_set(["ab", "ef", "de", "abc", "cd", "de", "acef"])) [['b', 'd', 'f']] """ # map elements to indices tgt = sorted(union(ss)) n = len(tgt) m = dict((x, i) for (i, x) in enumerate(tgt)) # X is a dict of sets (of indices) X = dict() for (i, s) in enumerate(ss): X[i] = set(m[x] for x in s) # Y shows which sets elements belong to Y = list(list() for _ in irange(n)) for (i, s) in enumerate(ss): for x in s: Y[m[x]].append(i) # find exact covers using algorithm X for rs in algorithmX(X, Y, list()): # turn the indices back to elements yield sorted(tgt[i] for i in rs) # 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 algorithm X 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 # a multiset variation on Knuth's Algorithm X def _mcover(m, tgt, X, ss, reject): # 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): # is this sequence acceptable? ss_ = ss + [n] if reject and reject(ss_): continue # update the target s = m[n] (tgt_, rs) = tgt.differences(s) if rs: 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, X[k].difference(discard)) for k in tgt_.keys()) if all(X_.values()): #yield from _mcover(m, tgt_, X_, ss_, reject) #[Python 3] for z in _mcover(m, tgt_, X_, ss_, reject): yield z #[Python 2] # 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. """ # 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 tgt.issuperset(s): 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(), reject) ############################################################################### # 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 __str__(self): return map2str((k, v) for (k, v) in self) 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__ # create a record from a dict @classmethod def from_dict(cls, d): return cls(**d) # 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)) i = 0 # count iterations while b - a > t: i += 1 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, fv) = ((x1, f1) if f1 < f2 else (x2, f2)) return Record(v=v, fv=fv, t=t, i=i) # 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_min(f, a, b, t=1e-9, m=None): """ find a 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: 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 (= f(v)) t = the tolerance used >>> r = find_min(lambda x: sq(x - 2), 0.0, 10.0) >>> round(r.v, 6) 2.0 """ return gss_minimiser(f, a, b, t=t, m=m) def find_max(f, a, b, t=1e-9): """ find a 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 (= f(v)) t = the tolerance used >>> r = find_max(lambda x: 9 - sq(x - 2), 0.0, 10.0) >>> round(r.v, 6) 2.0 """ r = find_min(f, a, b, t=t, m=neg) r.fv = neg(r.fv) return r # we can also use the minimiser to find roots # there are more rapidly converging root finding algorithms, # but unless f() is very expensive to call this will suffice def find_zero(f, a, b, t=1e-9, ft=1e-6): """ find a 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 (= f(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 """ # find a minimum of the absolute value r = find_min(f, a, b, t, m=abs) # and check the function value is close enough to 0 if ft < abs(r.fv): raise ValueError("value not found") # try a smaller t, or a larger ft r.ft = ft return r def find_value(f, v, a, b, t=1e-9, ft=1e-6): """ find a 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 (= f(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 """ fn = (f if v == 0 else (lambda x: f(x) - v)) r = find_zero(fn, a, b, t, ft) if v != 0: r.fv += v return r def find_values(f, v, a, b, t=1e-9, ft=1e-6, rt=1e-3): """ find all values x where f(x) = v for x in [a, b]. see: find_value() additional arguments: rt = minimum separation between values returned. """ # regions to check rs = [(a, b)] while rs: (a, b) = rs.pop(0) # look for a solution in this region try: r = find_value(f, v, a, b, t, ft) except ValueError: continue # return this value yield r # and consider the regions before and after this solution rs.extend([(a, r.v - rt), (r.v + rt, b)]) # can sides a, b, c form a triangle? @static( area=lambda v: (None if v < 0 else 0.25 * sqrt(v)), iarea=lambda v: (None if v < 0 else div(is_square(catch(as_int, v)), 4)), ) def is_triangle(a, b, c, fn=gt(0), validate=0): """ check it is possible to make a triangle with sides of length , , . by default the lengths are allowed to be 0 or negative, and the triangle formed uses the absolute value of the lengths. setting 'validate=1' will throw a ValueError for non-positive side lengths. the area of the triangle can be returned by setting: 'fn=is_triangle.area' when the area of the triangle will be returned, or None for impossible configurations. (degenerate triangles will return a zero area). for integer only areas: 'fn=is_triangle.iarea' can be used, when None will also be returned for non-integer areas. """ if validate and not (a > 0 and b > 0 and c > 0): raise ValueError("is_triangle: invalid side length ({a}, {b}, {c})".format(a=a, b=b, c=c)) # v = (4A)^2 where A is the area of the triangle v = (a + b + c) * (a + b - c) * (a + c - b) * (b + c - a) return fn(v) triangle_area = partial(is_triangle, fn=is_triangle.area) # return area as float (approximation) triangle_iarea = partial(is_triangle, fn=is_triangle.iarea) # return exact integer only areas # 2D geometry: # a point (P2) is represented by (x: Real, y: Real) # circle (Circle) is represented by (centre: P2, radius: Real) P2 = namedtuple('P2', 'x y') Circle = namedtuple('Circle', 'centre radius') def triangle_point(b, a, c, div=fdiv, sqrt=sqrt): """ find the position of the upper vertex of a triangle with base and other two sides and . the returned value is a point P = (x, y) with non-negative y value (height) where the base extends from O = (0, 0) to Q = (b, 0) (so |OQ| = b) and |OP| = a, |PQ| = c. see also: triangle_height() to just get the height of the triangle (y). """ (x, y) = (div(a*a + b*b - c*c, 2*b), None) if x is not None: y2 = a*a - x*x if y2 >= 0: y = sqrt(y2) else: x = None return P2(x, y) def triangle_height(b, a, c, div=fdiv, sqrt=sqrt): return triangle_point(b, a, c, div=div, sqrt=sqrt).y def triangle_iheight(b, a, c, div=rdiv, sqrt=(lambda x: is_square(as_int(x, include="0+", default=None)))): return triangle_point(b, a, c, div=div, sqrt=sqrt).y def triangle_circumcircle(A, B, C): """ find the circumcentre and circumradius of the triangle with vertices at points A, B, C. the returned value is a named tuple: (centre=(x, y), radius=r) """ ((Ax, Ay), (Bx, By), (Cx, Cy)) = (A, B, C) (A2, B2, C2) = (Ax*Ax + Ay*Ay, Bx*Bx + By*By, Cx*Cx + Cy*Cy) D = 2 * (Ax*(By - Cy) + Bx*(Cy - Ay) + Cx*(Ay - By)) # circumcentre U = P2( fdiv(A2*(By - Cy) + B2*(Cy - Ay) + C2*(Ay - By), D), fdiv(A2*(Cx - Bx) + B2*(Ax - Cx) + C2*(Bx - Ax), D) ) return Circle(U, point_distance(U, A)) def triangle_incircle(A, B, C): ((Ax, Ay), (Bx, By), (Cx, Cy)) = (A, B, C) (a, b, c) = (point_distance(B, C), point_distance(A, C), point_distance(A, B)) p = a + b + c # incentre I = P2(fdiv(a*Ax + b*Bx + c*Cx, p), fdiv(a*Ay + b*By + c*Cy, p)) return Circle(I, point_distance(I, A)) def point_distance(p1, p2): """ calculate the straight line distance between points p1 (= (x1, y1)) and p2 (= (x2, y2)) >>> point_distance((1, 2), (4, 6)) 5.0 """ ((x1, y1), (x2, y2)) = (p1, p2) return hypot(x1 - x2, y1 - y2) point_dist = point_distance def line_slope_intercept(p1, p2, div=fdiv): """ find the slope and intercept of a line defined by points p1 (= (x1, y1)) and p2 (= (x2, y2)). for non-vertical lines a pair (, ) is returned. for vertical lines (i.e. x = c), then (inf, c) is returned. >>> line_slope_intercept((0, 1), (2, 2)) (0.5, 1.0) >>> line_slope_intercept((1, 1), (1, 3)) (inf, 1) """ ((x1, y1), (x2, y2)) = (p1, p2) if x1 == x2: return (inf, x1) m = div(y2 - y1, x2 - x1) c = y1 - m * x1 return (m, c) 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)) q2 = (div(x - x3, dx43) if abs(dx43) > abs(dy43) else 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=P2(x, y), x=x, y=y, q1=q1, q2=q2) # t=0 -> p1; t=1 -> p2 def line_param(p1, p2, t=None): """ return a function that returns points on a line joining p1 and p2. if t = 0 then p1 is returned, and if t = 1 then p2 is returned. if t is specified then the function will be called with parameter t. """ ((x1, y1), (x2, y2)) = (p1, p2) f = lambda t, xd=x2 - x1, yd=y2 - y1: P2(x1 + t * xd, y1 + t * yd) return (f if t is None else f(t)) # 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 = div(x1 + x2 + y1 + y2, 2) return (P2(s - y1, s - x2), P2(s - y2, s - x1)) # return the closest distance between a line that passes through p1 and p2 # and the point p0 def line_distance(p1, p2, p0=(0, 0)): """ Return the minimum distance between the point p0 (= (x0, y0)) and a line that passes through points p1 (= (x1, y1)) and p2 (= (x2, y2)). If p0 is not specified (0, 0) is used. """ ((x1, y1), (x2, y2), (x0, y0)) = (p1, p2, p0) (xd, yd) = (x2 - x1, y2 - y1) return fdiv(abs(xd * (y1 - y0) - yd * (x1 - x0)), hypot(xd, yd)) # find intersection points of circle # with the line defined by points , # may return 0, 1, 2 points def circle_intersect_line(c, p1, p2): (((x0, y0), r), (x1, y1), (x2, y2)) = (c, p1, p2) (xd, yd) = (x2 - x1, y2 - y1) # construct a polynomial for the intersection f = sq(Polynomial([x1 - x0, xd])) + sq(Polynomial([y1 - y0, yd])) ts = f.roots(sq(r), domain='F') return list(P2(x1 + t * xd, y1 + t * yd) for t in ts) # calculate the points where two circles intersect def circle_intersect_circle(c1, c2): ((o1, r1), (o2, r2)) = (c1, c2) # calculate the distance between the centres d = point_dist(o1, o2) if d == 0 or r1 <= 0 or r2 <= 0: raise ValueError("circle_intersect_circle: invalid circles") # the fraction of the intersect points along the line joining the centres fn = line_param(o1, o2) t = 0.5 + fdiv(sq(r1) - sq(r2), 2.0 * sq(d)) # find the perpendicular bisector of the line at t (p1, p2) = line_bisect(o1, fn(2*t)) # return the intersections with the first circle return circle_intersect_line(c1, p1, p2) # parameterised circle t = -1 to +1 for a full circle # 0 = 3 o'clock; 0 -> 1 = anticlockwise to 9 o'clock; 0 -> -1 = clockwise to 9 o'clock def circle_param(c, t=None): ((x, y), r) = c f = lambda t, x=x, y=y, r=r: P2(x + r * cos_ht(t, 1), y + r * sin_ht(t, 1)) return (f if t is None else f(t)) # find parameter for point

on circle centre radius def circle_param_t(c, p, t=1e-9): (o, r) = c if not (abs(point_dist(o, p) - r) < t): return None ((x0, y0), (x, y)) = (o, p) t = acos_ht(fdiv(x - x0, r), 1) return min(t, -t, key=(lambda t: abs(point_dist(circle_param(c, t), p)))) def polygon_area(ps, m=0.5, sum=math.fsum): """ calculate the (signed) area of a simple (non-intersecting) polygon from a sequence of (x, y) points specifying the vertices >>> polygon_area([(0, 0), (1, 0), (1, 1), (0, 1)]) 1.0 """ return m * sum(x1 * y2 - x2 * y1 for ((x1, y1), (x2, y2)) in tuples(ps, 2, circular=1)) # cache the areas of unit 2-,3-,4-,6-,8-gons @static(cache={ 2: 0.0, 3: 0.4330127018922193, 4: 1.0, 6: 2.598076211353316, 8: 4.82842712474619 }) def polygon_area_regular(n, x=1): """ calculate the area of a regular -gon with side """ cache = polygon_area_regular.cache # calculate the area of a unit polygon a = cache.get(n) if a is None: a = fdiv(0.25 * n, math.tan(fdiv(pi, n))) cache[n] = a if x == 1: return a return a * x * x ############################################################################### # Roman Numerals # the following is good for numbers in the range 1 - 4999 # values for letters in roman numerals _romans = { 'M': 1000, 'D': 500, 'C': 100, 'L': 50, 'X': 10, 'V': 5, 'I': 1, } # place values for roman numerals _int2roman = ( ("", "I", "II", "III", "IV", "V", "VI", "VII", "VIII", "IX"), # units ("", "X", "XX", "XXX", "XL", "L", "LX", "LXX", "LXXX", "XC"), # tens ("", "C", "CC", "CCC", "CD", "D", "DC", "DCC", "DCCC", "CM"), # hundreds ("", "M", "MM", "MMM", "MMMM"), # thousands ) def int2roman(n, validate=0): """ return a representation of an integer (from 1 to 4999) as a Roman Numeral >>> int2roman(4) 'IV' >>> int2roman(1999) 'MCMXCIX' """ if validate: n = as_int(n) if not (0 < n < 5000): raise ValueError("integer out of range: {n}".format(n=n)) return str.join('', list(_int2roman[i][x] for (i, x) in enumerate(nsplitter(n)))[::-1]) def roman2int(r, strict=1): """ return the integer value of a Roman Numeral . if 'strict' is set (the default), then the function only accepts Roman Numerals in "standard" form. if 'strict' is not set, then the function is generous with what it accepts as a Roman Numeral. >>> roman2int('IV') 4 >>> roman2int('IIII') 4 >>> roman2int('MCMXCIX') 1999 >>> roman2int('XLXV', strict=0) 55 """ r = str(r).upper() if r == 'IIII': return 4 # group letters into clumps, and value the clumps try: vs = list(_romans[s[0]] * len(s) for s in clump(r)) except KeyError: raise ValueError("invalid Roman numeral: {r!r}".format(r=r)) # evaluate the list n = i = 0 while True: if i + 1 < len(vs): (a, b) = (vs[i], vs[i + 1]) if a < b: # subtractive form n += b - a i += 2 else: # additive form n += a i += 1 else: if i < len(vs): n += vs[i] break if strict and int2roman(n) != r: raise ValueError("invalid roman numeral: {r!r}".format(r=r)) return n def is_roman(r): """ check if a Roman Numeral is valid (i.e. in "standard" form). >>> is_roman('IV') 4 >>> is_roman('IIII') 4 >>> is_roman('XIVI') 0 """ r = str(r).upper() if r == 'IIII': return 4 return catch(roman2int, r, strict=1) or 0 # set the default base (radix) def base(*args): if len(args) > 1: raise TypeError(str.format("base: invalid base: {args!r}", args=args)) if args: n = as_int((args[0] or 10), include="+") if n < 2: raise ValueError(str.format("base: invalid base: {n}", n=n)) enigma.radix = (args[0] or 10) return enigma.radix # digits = (default) digits for use in converting bases @static(digits=str_digit + str_upper) 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. """ if len(args) > 1: raise TypeError(str.format("base_digits: invalid base digits: {args!r}", args=args)) if args: base_digits.digits = (args[0] or (str_digit + str_upper)) return base_digits.digits def int2base(i, base=0, width=None, pad=None, group=None, sep=",", neg='-', digits=None, validate=0): """ 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, using standard digits 0-9 and then letters A-Z, but the parameter can be specified to give the symbols for larger bases. And if there are more digits in the specified base then there are available symbols, then the returned string will be of the form "{::...}" where is the digit value expressed in decimal (using digits 0-9). alias: int2str(). see also: gmpy2.digits() >>> 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' """ if base == 0: base = radix if base < 2: raise ValueError("invalid base {base!r}".format(base=base)) if digits is None: digits = base_digits() i = (as_int(i) if validate else int(i)) (p, r) = ('', None) if i < 0: (p, i) = (neg, -i) # if there aren't enough digits switch to {::...} format if len(digits) < base: ds = nsplit(i, base=base, fn=list) # pad left (or right for negative values) with zero digit if width and width > 0: while len(ds) < width: ds.insert(0, 0) elif width and width < 0: while len(ds) < -width: ds.append(0) r = join(ds, sep=':', enc="{}") width = group = None # group is just ignored for now 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: if group < 0: r = join((join(x) for x in chunk(r, -group)), sep=sep) if group > 0: r = rev(join((join(x) for x in chunk(rev(r), group)), sep=rev(sep))) return (p + r if p else r) int2str = int2base def base2int(s, base=0, strip=0, neg="-~", digits=None): """ convert a string representation of an integer in the specified base to an integer. if is set, then invalid characters in the conversion are ignored. alias: str2int(). >>> base2int('-42') -42 >>> base2int('xyzzy', base=3, digits='zyx') 190 >>> base2int('HELLO', base=36) 29234652 """ if base == 0: base = radix if base < 2: raise ValueError("invalid base {base!r}".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 i = fneg = 0 if s and s[0] in neg: # - (or ~) indicate a negative integer fneg ^= 1 i += 1 n = f = 0 for k in range(i, len(s)): d = s[k] try: v = digits.index(d) except ValueError as e: if strip: continue e.args = ("invalid digit {d!r} for base {base} in {s!r}".format(base=base, s=s, d=d),) raise n *= base n += v f = 1 if f == 0: return None # if no characters processed return (-n if fneg else n) str2int = base2int 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) # there is a 'num2words' library that handles different languages _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: pass raise ValueError(str.format('Number too large (scale={scale!r})', 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 = (ndigits(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; '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(str.format("int2words: lang='{lang}' not implemented", lang=lang)) 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=0, 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 """ if base == 0: base = radix 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 (r if s == 1 else -r) # convert a sequence to a string: "(a, b, c)" def seq2str(s, enc="()", sep=", ", sort=0, rev=0, fn=str): """ convert a sequence to a string suitable for output. alias: join() >>> seq2str([2, 1, 3]) '(2, 1, 3)' >>> seq2str([2, 1, 3], sort=1) '(1, 2, 3)' >>> seq2str(first(primes, 10)) '(2, 3, 5, 7, 11, 13, 17, 19, 23, 29)' >>> seq2str(map(seq2str, prime_factor(factorial(15) - 1))) '((17, 1), (31, 2), (53, 1), (1510259, 1))' """ return join(s, sep=sep, enc=enc, sort=sort, rev=rev, fn=fn) # convert a map to a string: "(a=1, b=2, c=3)" def map2str(m, sort=1, enc="()", sep=", ", arr="="): """ convert a map to a string suitable 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) # is it a dict? if isinstance(m, Mapping): return join((concat(k, arr, m[k]) for k in fn(m.keys())), sep=sep, enc=enc) # (k, v) pairs return join((concat(k, arr, v) for (k, v) in fn(m)), sep=sep, enc=enc) ############################################################################### # 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) for example: x = Delay(expensive, 1) x.evaluated --> False x.value (or x()) --> returns expensive(1) x.evaluated --> True x.value (or x()) --> returns expensive(1) again, without re-evaluating it x.reset() x.evaluated --> False x.value (or x()) --> 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 = call(self.fn, self.args, self.kw) self.evaluated = True return self.value def __call__(self): if not self.evaluated: self.evaluate() return self.value def reset(self): del self.value self.evaluated = False def __getattr__(self, key): if key == 'value': return self.evaluate() 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(fn=add) >>> 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 v0 = self.value self.value = (self.fn1(v) if v0 is None else self.fn(v0, 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, seq): """accumulate values from sequence """ for v in seq: self.accumulate(v) return self def accumulate_data_from(self, seq, value=0, data=1): """ accumulate values and data from sequence . , can be an index into elements from or a function to extract the appropriate value from an element. """ if not callable(value): value = item(value) if not callable(data): data = item(data) for x in seq: 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] ############################################################################### # sequence iterator/enumerator # acts like enumerate(), but has an additional 'count' attribute class Enumerator(object): """ iterate through a sequence, counting the number of items returned >>> ss = Enumerator(primes.between(10, 99)) >>> list(ss.iter()) [11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97] >>> ss.count 21 >>> ss = Enumerator(empty) >>> list(ss) [] >>> ss.count 0 """ def __init__(self, seq, count=0): self.seq = seq self.count = count # count the number of items returned def iter(self, skip=0): """iterate through the sequence, counting the number of items returned""" for (i, x) in seq_items(self.seq, skip): self.count += 1 yield x __iter__ = iter def enumerate(self, skip=0, offset=0): """enumerate the sequence, counting the number of items returned""" for (i, x) in seq_items(self.seq, skip): self.count += 1 yield (i + offset, x) def output(self, skip=0, offset=0, fmt="{count}: {element}"): """output the contents of the enumerator""" for (count, element) in self.enumerate(skip=skip, offset=offset): printf(fmt) ############################################################################### # non-empty closed intervals of integers: # represent an interval by it's min and max values class Interval(object): def __init__(self, mn=None, mx=None): self.mn = 0 # lower limit self.mx = inf # upper limit self.update(mn, mx) def __repr__(self): if self.mn == self.mx: return sprintf("<{self.mn}>") # singleton return sprintf("<{self.mn},{self.mx}>") # update the min/max values def update(self, mn=None, mx=None): if mn is not None and mn > self.mn: self.mn = mn if mx is not None and mx < self.mx: self.mx = mx if self.mn > self.mx: raise ValueError(sprintf("Bad Interval {self}")) return self # size of an interval def size(self): return self.mx - self.mn + 1 # iterate through the values in an interval def __iter__(self): for i in irange(self.mn, self.mx): yield i # does an interval contain a value? def contains(self, x): return self.mn <= x <= self.mx # return a copy of this interval def copy(self): return Interval(mn=self.mn, mx=self.mx) # an interval that is the sum of this interval and another def sum(self, *others): return Interval(mn=self.mn + sum(i.mn for i in others), mx=self.mx + sum(i.mx for i in others)) # an interval that is the intersection of this interval and another def intersect(self, other): return Interval(mn=max(self.mn, other.mn), mx=min(self.mx, other.mx)) ############################################################################### # Routines for dealing with polynomials # represent polynomial a + b.x + c.x^2 + d.x^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 # parse a string as a polynomial def poly_from_str(s, var='x'): if s == var: return [0, 1] if s == '1': return [1] if s == '0': return [0] d = defaultdict(int) sgn = +1 # split string into segments for x in re.split(r'\s*([+-])\s*', respace(s, sep='')): if x == '': continue elif x == '-': sgn = -1 elif x == '+': sgn = +1 else: i = x.find(var) if i == -1: # constant term (= x^0) d[0] += sgn * number(x) else: # variable found, parse coefficient and exponent c = number(x[:i]) if c is None: c = 1 e = number(x[i:]) if e is None: e = 1 d[e] += sgn * c return poly_from_pairs(d.items()) # 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 # make a polynomial with the given roots def poly_from_roots(rs): return poly_multiply([-r, 1] for r in rs) # 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 try: # are coefficients already integers? p = list(as_int(x) for x in p) except ValueError: # assume rationals 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 """ if domain not in "QZ": raise ValueError(str.format("invalid domain: {domain!r}", domain=domain)) 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': x = as_int(x, default=None) if x is None: continue 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

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 # TODO: consider allowing an "offset" (or "shift") parameter that multiplies each coefficient # by x^n, which would allow us to do some handling of polynomials with negative coefficients # e.g.: a.x^-2 + b.x^-1 + c + d.x + e.x^2 = (a + b.x + c.x^2 + d.x^3 + e.x^4) . x^-2 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, ...] coefficients are intended to be integers or rational numbers. """ def __init__(self, obj, **kw): if isinstance(obj, basestring): obj = poly_from_str(obj, var=kw.get('var', 'x')) list.__init__(self, obj) 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 function(self): "return a Python function that calls the polynomial" return partial(poly_value, self) 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 quadratic_roots(self, v=0, domain='Q', include="+-0", F=None): p = (self - v if v else self) if p.degree() > 2: raise ValueError("Polynomial.quadratic_roots: degree too large") return quadratic(p.coeff(2), p.coeff(1), p.coeff(0), domain=domain, include=include, F=F) def cubic_roots(self, v=0, domain='Q', include="+-0", F=None): p = (self - v if v else self) if p.degree() > 3: raise ValueError("Polynomial.cubic_roots: degree too large") return cubic(p.coeff(3), p.coeff(2), p.coeff(1), p.coeff(0), domain=domain, include=include, F=F) def rational_roots(self, v=0, domain='Q', include="+-0", F=None): "find rational roots of the polynomial = v" p = (self - v if v else self) return poly_rational_roots(p, domain=domain, include=include, F=F) def roots(self, v=0, domain='Q', include="+-0", F=None): """ find all roots of the polynomial = v that satisfy the specified requirements. domain = 'Z' (integer roots), 'Q' (rational roots), 'F' (real (float) roots), 'C' ( complex roots). for non-complex roots, can be specified to include just roots that some combination of zero = '0', positive = '+', negative = '-'. currently the routine cannot find non-integer/non-rational roots for irreducible polynomials of degree 4 or more. """ p = (self - v if v else self) k = p.degree() if k < 3: return quadratic(p.coeff(2), p.coeff(1), p.coeff(0), domain=domain, include=include, F=F) if k < 4: return cubic(p.coeff(3), p.coeff(2), p.coeff(1), p.coeff(0), domain=domain, include=include, F=F) if domain in 'QZ': return p.rational_roots(domain=domain, include=include, F=F) raise ValueError("Polynomial.roots: polynomial too complicated, maybe try find_roots()") 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)) def find_roots(self, v=0, domain='Q', F=None, div=None, warn=1): "try to find roots of a polynomial by factorisation" p = (self - v if v else self) for (f, n) in p.factor(F=F, div=div): d = f.degree() if d < 1: continue # currently we can only deal non-rational roots in (up to) cubic factors if d > 3: if domain in 'FC' and warn: printf("WARNING: Polynomial.find_roots: ignoring poly factor {f}") # TODO: consider using mpmath.polyroots() here continue for x in f.cubic_roots(domain=domain, F=F): yield x @classmethod def from_pairs(cls, ps): return cls(poly_from_pairs(ps)) @classmethod def from_roots(cls, rs): return cls(poly_from_roots(rs)) @classmethod def from_str(cls, s, var='x'): return cls(poly_from_str(s, var=var)) @classmethod def unit(cls): return cls(poly_unit) # -> 1 @classmethod def zero(cls): return cls(poly_zero) # -> 0 @classmethod def linear(cls, a=1, b=0): return cls([b, a]) # -> a.x + b @classmethod def quadratic(cls, a=1, b=0, c=0): return cls([c, b, a]) # -> a.x^2 + b.x + c @classmethod def cubic(cls, a=1, b=0, c=0, d=0): return cls([d, c, b, a]) # -> a.x^3 + b.x^2 + c.x + d # 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] >>> prms = _PrimeSieveE6(1000000) >>> prms.is_prime(10001) False >>> 10007 in prms True >>> sum(prms) == 37550402023 True >>> list(prms.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 expand the sieve to the required size self.expand(n) def __repr__(self): return self.__class__.__name__ + '(max=' + repr(self.max) + ')' def expand(self, n): """ expand 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) # generate k-subsets of primes from the sieve (in the interval [a, b]) def subsets(self, k, a=0, b=inf): # consider increasing values for the largest prime for p in self.irange(a, b): if k == 1: yield (p,) elif k > 1: # generate the smaller primes for ps in self.subsets(k - 1, a=a, b=p - 1): yield ps + (p,) tuples = subsets # best to use subsets() # prime test (may throw IndexError if n is too large) def contains(self, n, validate=0): """ check if the number is a prime in the sieve. this is a constant time lookup in the sieve, and may throw IndexError for numbers larger than the sieve. """ if n is None: return None if validate: n = as_int(n, include="0+") if n < 5: return (n == 2 or n == 3) # 2, 3 -> T; 0, 1, 4 -> F r = n % 6 if r != 1 and r != 5: return False # (n % 6) != (1, 5) -> F return bool(self.sieve[n // 3]) # alternative "is_prime" test by searching the sieve for prime factors def is_prime(self, n, validate=0): """ check if the number is prime, either if it is contained in the sieve or by searching for prime factors that are contained in the sieve. this can test numbers up to the square of the size of the sieve (which will be expanded if necessary). this is the function used to implement [[ n in prms ]]. """ if n is None: return None if validate: n = as_int(n, include="0+") if n < 5: return (n == 2 or n == 3) r = n % 6 if r != 1 and r != 5: return False # for numbers less then the sieve size, we can just look them up if n <= self.max: return bool(self.sieve[n // 3]) # otherwise we look for prime factors r = isqrt(n) if self.expandable: self.expand(r) for p in self.irange(2, r): if n % p == 0: return False return True prime = is_prime # allows use of "in" __contains__ = is_prime def is_composite(self, n, validate=0): if n is None: return None if validate: n = as_int(n, include="0+") if n < 4: return False return not self.is_prime(n) # 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: while 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 the interval [a, b] """ if self.expandable: self.expand(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, expand=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. However, if 'expand' is set, then the sieve will be expanded to include primes up to isqrt(n) before factorisation is performed. """ if n == 1: return () if expand: self.expand(isqrt(n)) return prime_factor_h(n, self, end=end, nf=mr, mr=mr, mrr=mrr) def prime_factors(self, n, end=None, mr=0, mrr=0, expand=0): for (p, e) in self.prime_factor(n, end=end, mr=mr, mrr=mrr, expand=expand): for _ in range(e): yield p # 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, expand=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=partial(self.prime_factor, end=end, mr=mr, mrr=mrr, expand=expand)) def divisors(self, n, end=None, mr=0, mrr=0, expand=0): return divisors(n, fn=partial(self.prime_factor, end=end, mr=mr, mrr=mrr, expand=expand)) def divisors_pairs(self, n, end=None, mr=0, mrr=0, expand=0): return divisors_pairs(n, fn=partial(self.prime_factor, end=end, mr=mr, mrr=mrr, expand=expand)) def tau(self, n, end=None, mr=0, mrr=0, expand=0): return tau(n, fn=partial(self.prime_factor, end=end, mr=mr, mrr=mrr, expand=expand)) def is_square_free(self, n, end=None, mr=0, mrr=0, expand=0): return is_square_free(n, fn=partial(self.prime_factor, end=end, mr=mr, mrr=mrr, expand=expand)) def factorisations(self, n, end=None, mr=0, mrr=0, expand=0): return factorisations(n, fn=partial(self.divisors, end=end, mr=mr, mrr=mrr, expand=expand)) # 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): >>> prms = _PrimeSieveE6X(1000) >>> first(prms, 1, skip=999) [7919] We can then find the one millionth prime and the generator will expand as necessary: >>> first(prms, 1, skip=999999) [15485863] We can see what the current maximum number considered is: >>> prms.max 16384000 And can quickly test for primality up to this value: >>> 1000003 in prms True But we can use the sieve to test for larger primes (and it will be expanded as necessary): >>> 17000023 in prms True """ 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 expand(self, n=None): """ expand 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.expand(self, n) return self # for backwards compatibility extend = expand # 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.expand(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: >>> prms = Primes(50) >>> prms.contents() [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47] >>> sum(prms) 328 >>> 39 in prms False The collection can be extended manually to a new upper limit: >>> prms.expand(100) >>> sum(prms) 1060 >>> 97 in prms True but it doesn't automatically expand. If we want an automatically expanding version, we can set the 'expandable' flag to True. >>> prms = Primes(50, expandable=1) We can find out the current size and contents of the sieve: >>> prms.max 50 >>> prms.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(prms.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(prms, 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) 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: deepcopy = lazy_import('copy.deepcopy') # 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 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 ############################################################################### # output sums def output_mul(a, b, base=0, rev=0, pre='', start=None, end=None, sep=''): """output * as a long multiplication sum""" if base == 0: base = radix (a, b) = (as_int(x, include='0+') for x in (a, b)) c = a * b #printf("{a} * {b} = {c}") (ka, kb, kc) = (ndigits(x, base=base) for x in (a, b, c)) k = max(ka, kb, kc) # format a number fmt = lambda n, base=base, pad=' ', width=k: int2base(n, base=base, pad=pad, width=width, group=1, sep=sep) (fa, fb, fc) = (fmt(a), fmt(b), fmt(c)) w = max(len(fa), len(fb), len(fc)) # output the multiplication if start is not None: printf("{start}") printf("{pre}{fa}") printf("{pre}{fb} {sep}*") printf("{pre}{x}", x='-' * w) # if b != 0: ... # output the intermediates xs = list() for (i, d) in enumerate(nsplit(b, base=base, reverse=1)): p = a * d if p == 0: continue # skip 0 products #printf("{pre}{p}{x} = {d} * {a}", p=fmt(p, width=k - i), x=' ' * i) xs.append(fmt(p, width=k - i)) if rev: xs.reverse() for x in xs: printf("{pre}{x}") # output the result printf("{pre}{x}", x='-' * w) printf("{pre}{fc}") printf("{pre}{x}", x='=' * w) if end is not None: printf("{end}") # generate an underline string for strings in def underline(ss, ul='-', spc=' '): # all strings should be the same length n = xl = xr = None for s in ss: if s is None: continue # skip None if n is None: n = len(s) (xl, xr) = (n - 1, 0) xl = peek((i for i in irange(0, xl - 1) if s[i] != spc), default=xl) xr = peek((i for i in irange(n - 1, xr + 1, step=-1) if s[i] != spc), default=xr) return (spc * xl) + (ul * (xr - xl + 1)) + (spc * (n - xr - 1)) def output_div(a, b, rem=0, base=0, pre='', start=None, end=None, sep=''): """ output / as a long division sum if is set then the remainder will always be printed (even if it is zero). """ if base == 0: base = radix (c, r) = divmod(a, b) #printf("{a} / {b} = {c} (rem {r})") (ka, kb, kc) = (ndigits(x, base=base) for x in (a, b, c)) # format a number fmt = partial(int2base, base=base, pad=' ', group=1, sep=sep) (fa, fb, fc) = (fmt(a), fmt(b), fmt(c)) # output the division if start is not None: printf("{start}") printf("{pre}{z}{fc}", z=' ' * (len(fb) + len(fa) - len(fc) + 5)) printf("{pre}{z} ---{zz}", z=' ' * len(fb), zz='-' * len(fa)) printf("{pre}{fb} ) {fa}") (ds, p, w, s, z) = (nsplit(a, base=base, fn=list), 0, 0, None, ' ' * (len(fb) + 5)) while ds: p = base * p + ds.pop(0) q = floor(p, b) w += 1 if q > 0: (fp, fq) = (fmt(p, width=w), fmt(q, width=w)) if s: printf("{pre}{z}{s}") printf("{pre}{z}{fp}") printf("{pre}{z}{fq}") (s, p) = (underline([fp, fq]), p - q) if p or rem: if s: printf("{pre}{z}{s}") printf("{pre}{z}{p} (rem)", p=fmt(p, width=w)) s = underline([fp, fq]) printf("{pre}{z}{s}", s=s.replace('-', '=')) if end is not None: printf("{end}") def output_sqrx(x, r=None, rem=0, base=0, pre='', start=None, end=None, sep=''): """output the extraction of the square root of """ if base == 0: base = radix if start is not None: printf("{start}") if r is None: r = isqrt(x) n = ndigits(x, base=base) fmt = partial(int2base, base=base, pad=' ', group=1, sep=sep) fx = fmt(x) fr = fmt(r, sep=sep + ' ' + sep) printf("{pre} {z}{fr}", z=' ' * (len(fx) - len(fr))) printf("{pre} -{z}", z='-' * len(fx)) printf("{pre}\\/ {fx}") # chop digits of x in pairs xds = nsplit(x, base=base*base) rds = nsplit(r, base=base) w = (1 if xds[0] < base else 2) r = z = 0 fz = None sqrx.base = base for (i, (xd, d)) in enumerate(zip(xds, rds)): z = z*base*base + xd if d > 0: fz = fmt(z, width=w) if i > 0: printf("{pre} {fz}") p = d * (2 * base * r + d) # = sqrx(r, d) z -= p if d > 0: fp = fmt(p, width=w) printf("{pre} {fp}") s = underline([fz, fp], ('=' if (z == 0 and i + 1 == len(rds) and not rem) else '-')) printf("{pre} {s}") w += 2 r = r*base + d if z or rem: printf("{pre} {z} (rem)", z=fmt(z, width=w - 2)) printf("{pre} {s}", s=underline([fz], '=')) if end is not None: printf("{end}") ############################################################################### # Substituted Sum Solver # NOTE: SubstitutedExpression.split_sum() can do all this, and more, faster. # 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=0): """ 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) """ if base == 0: base = radix # check there aren't too many letters (we could issue a warning) words = list(terms) words.append(result) if len(union(words)) > base: return # 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 = { 0: set(w[0] for w in words if len(w) > 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) if not all(len(t) == n for t in ts): raise ValueError("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): """ Note: See SubstitutedExpression.split_sum() for a more powerful solver. 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=0, 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 = (radix if base == 0 else 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=0, digits=None, l2d=None, d2i=None): """ solve a sequence of substituted sum problems. sums are specified as a sequence of: (, , ..., ) """ # are we done? if base == 0: base = radix 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=0, digits=None, l2d=None, d2i=None): if base == 0: base = radix 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'] = numbers(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 numbers(ds, csep=_split_sep): 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 # solve the alphametic puzzle + ... + = # where one of the symbols given is incorrect # is a list of the indices of suspect terms # (see: Puzzle 56 [ https://enigmaticcode.wordpress.com/2018/01/24/puzzle-56-addition/ ]) @classmethod def bungled_sum(cls, ts, ss=None, base=0): if base == 0: base = radix # sort out the arguments n = len(ts) if ss is None: ss = list(range(n)) else: if not all(s < n for s in ss): raise ValueError("invalid suspect index") # all letters in all terms letters = union(ts) if len(letters) < base: letters.add('?') # check the suspect terms for i in set(ss): t = ts[i] for (j, x) in enumerate(t): # choose a replacement letter for y in letters: if x == y: continue # create a puzzle for the modified sum terms = update(ts, [(i, update(t, [(j, y)]))]) if y == '?' and x not in union(terms): continue p = cls(terms[:-1], terms[-1], base=base) # solve the puzzle for (k, s) in enumerate(p.solve(verbose=0)): if k == 0: # indicate the incorrect symbol h = ' ' * (sum(len(t) + 3 for t in terms[:i]) + j) + x printf("\n{h}\n{p.text} / @[{i},{j}] {x} -> {y}") # output solutions p.output_solution(s) @classmethod def run_bungled_sum(cls, args): ts = ss = None opt = dict() fail = 0 for arg in args: k = v = None try: if arg.startswith('--'): (k, _, v) = arg.lstrip('-').partition('=') elif arg.startswith('-'): (k, v) = (arg[1], arg[2:]) if k is None: # plain arguments if ts is None: # the sum ts = re.split(r'[\s+\=]+', arg) else: # an index if ss is None: ss = list() ss.append(int(arg)) else: # options if k == 'b' or k == 'base': opt['base'] = int(v) else: fail = 1 except Exception as e: print(e) fail = 1 break if fail: printf("usage: {cls.__name__}.bungled_sum [--base=] \" + + ... = \" [ ...]") return -1 cls.bungled_sum(ts, ss, **opt) return 0 bungled_sum = SubstitutedSum.bungled_sum ############################################################################### # 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 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] # # TODO: consider having an inverse "--invalid", so X is odd: # --valid="1|3|5|7|9,X" [instead of: --invalid="0|2|4|6|8,X"] # find words in string def _find_words(s, r=1): words = set(re.findall(r'{(\w+?)}', s)) # re.UNICODE 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: # default is just _x _sym = lambda x: '_' + x # this will use ascii variables 'v_' #_sym = cached(lambda x: 'v_' + join((int2base(ord(c), base=16) for c in x), sep="_")) _set = partial(join, sep=", ", enc="{}") # 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) # NOTE: file.writelines does NOT include newline characters def writelines(fh, lines, sep=os.linesep, flush=1): for line in lines: fh.write(line) if sep: fh.write(sep) if flush: fh.flush() def _readlines(fh, strip=3, fn=None, st=bool): while True: x = fh.readline() if not x: break if strip & 2: x = x.partition('#')[0] # strip #-comments if strip & 1: x = x.strip() # strip whitespace if fn: x = fn(x) if st(x): yield x # read lines from @static(WHITESPACE=1, COMMENTS=2) # strip flags def readlines(fh, strip=3, fn=None, st=bool, first=None): """ read lines from file handle strip = 1 will strip trailing whitespace (including newlines) strip = 2 will strip #-comments strip = 3 does both (default) fn = function to manipulate (stripped) lines (default = None) st = check for (stripped, manipulated) lines (default = bool) first = stop after returning the first valid lines """ # just read the first few lines? if first is not None: return ifirst(_readlines(fh, strip=strip, fn=fn, st=st), first) else: return _readlines(fh, strip=strip, fn=fn, st=st) # 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(encl(re.escape(sep[1:]), '[]'), d, s) if maxsplit == 0: return s.split(d) elif maxsplit > 0: return s.split(d, maxsplit) else: return s.rsplit(d, -maxsplit) # check for explicit "{WORD}" parameters _is_explicit = lambda s: re.search(r'{\w+}', s) # was: [[ ('{' in s) ]] # 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 _is_explicit(s): return s return re.sub(encl(symbols, '[]') + '+', (lambda m: encl(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) def _expand_macros(s, macro, depth=20): if not s: return s # this leaves invalid macros alone fn = lambda k: macro.get(k[1:], k) # macro[k[1:]] # raises error for invalid macros while '@' in s or '$' in s: if depth == 0: raise ValueError("macro depth exceeded") s_ = s try: s = re.sub(r'([\@\$]\w+)', (lambda x: fn(x.group(1))), s) except KeyError as e: # [Python 3] raise ValueError(...) from None x = ValueError(str.format("invalid macro: {k!r}", k=e.args[0])) if hasattr(x, '__cause__'): x.__cause__ = None raise x if s == s_: break # stop if line remains unchanged depth -= 1 return s 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(), allow. 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=0, distinct=1, process=1, reorder=1, first=0, denest=0, sane=1, warn=0, opt='', 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, output=None, check=None, macro=None, env=None, code=None, decl=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, output=None, distinct=None, check=None, macro=None, env=None, code=None, process=None, reorder=None, first=None, denest=None, decl=None, sane=None, warn=None, opt=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) template - output template (default: derived from exprs) solution - output solution values (default: derived from exprs) header - output header (default: derived from exprs) output - provide custom output function (default: None) answer - an expression for the answer value (default: None) accumulate - accumulate answers using the specified object (default: None) 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) macro - macro expansions applied to expressions containing @macro (default: None) denest - work around CPython statically nested block limit decl - additional declarations used in functions generated when denest is enabled 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(): v = (get_default(k, s2d, l2d) if k == 's2d' else 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) vW = 1024 # output warnings # standard debug levels v0 = 0 # 0 = no output v1 = vH | vT | vA | vW # 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 | vW # 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 "new style" flags, e.g. "1+E" "1-T" "1-T+E" "9" "HTAEPIC" d = dict((k, getattr(self, 'v' + k)) for k in 'HTAEPICW01239') v = d['0'] try: op = '+' for k in n: if k in '+-': op = k continue x = d[k] if op == '+': v |= x elif op == '-': v &= ~x op = '+' return v except Exception: pass # sort out "old-style" numeric arguments if re.match(r'[\d\-\|\+\,]+', n): n = sum(numbers(n, csep='|+,')) # 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 solution = self.solution header = self.header distinct = self.distinct literal = self.literal code = self.code denest = self.denest decl = self.decl macro = self.macro process = self.process sane = self.sane warn = self.warn opt = self.opt verbose = self.verbose # sort out verbose argument verbose = self._verbose(verbose) if verbose & self.vW: if opt: xs = list(v for (k, v) in {'i': 'iterators', 's': 'sets'}.items() if k in opt) if xs: printf("WARNING: enabling experimental options: {xs}", xs=join(xs, sep=", ")) if sane == 0 : printf("WARNING: sanity checks disabled - good luck!") # base (use default radix if specified as 0) if base == 0: base = radix # the symbols to replace (for implicit expressions) if symbols is None: symbols = str_upper symbols = set(symbols) if s2d: symbols.update(s2d.keys()) if literal: symbols.update(literal) if solution: symbols.update(solution) symbols = join(sorted(symbols)) # code should be a sequence of strings if code and isinstance(code, basestring): code = [code] # 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] # replace any macros within expressions (and answer, template) if macro: exprs = list(_expand_macros(x, macro) for x in exprs) answer = _expand_macros(answer, macro) template = _expand_macros(template, macro) if code: code = list(_expand_macros(x, macro) for x in code) # now process the list (skipping duplicates) xs = list() for expr in uniq(exprs, fn=repr): 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 and not (_is_explicit(expr) or (isinstance(v, basestring) and _is_explicit(v))): expr = _fix_implicit(expr, symbols) if isinstance(v, basestring): v = _fix_implicit(v, symbols) # value is either an alphametic or a numeric literal if isinstance(v, basestring) and not (v[0] == '{' and v[-1] == '}'): 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/solution ( = 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(ts, fn=(lambda t: encl(t, '()')), 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 & self.vW: # 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 need to be filled out symbols = set().union(*words) # make sure all specified symbols are included if self.symbols: symbols.update(self.symbols) # 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 & self.vW 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 on words with multiple symbols 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) # we only need to construct words with more than 1 symbol words = set(w for w in words 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 = (join(symbols, sort=1) 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 if template is None: template = _template if solution is None: solution = join(diff(symbols, literal), sort=1) if header is None: header = _replace_words(template, identity) # complete list of symbols to determine if distinct: symbols.update(*distinct) if solution: symbols.update(solution) symbols = join(symbols, sort=1) # sort out negative value for denest, will not be enabled if running under PyPy if denest < 0: denest = (0 if _pypy else -denest) if denest == 1: denest = 50 # default denest value # update the processed values self.exprs = exprs self.base = base self.symbols = symbols self.digits = digits self.s2d = s2d self.answer = answer self.template = template self.solution = solution self.header = header self.distinct = distinct self.literal = literal self.code = code self.denest = denest self.decl = decl self.verbose = verbose self._words = words self._invalid = invalid self._idigits = idigits self._exprs = (exprs, xs, vs, ts, syms) self._fail = 0 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, verbose=None): 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 output = self.output reorder = self.reorder denest = self.denest decl = self.decl sane = self.sane warn = self.warn opt = self.opt if verbose is None: 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 if not symbols: raise ValueError("expr has no alphametic symbols ???") 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 & self.vW: 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 & self.vW: printf("[SubstitutedExpression: replacing ({t}) -> ({val} = {{{word}}})]", t=ts[i]) exprs[i] = (int2base(val, base=radix), word, 3) (xs[i], vs[i]) = (vs[i], xs[i]) # reorder the expressions into a more appropriate evaluation order if reorder: # at each stage choose 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) # if reorder is a list, do the specified indices (1-indexed) first if isinstance(reorder, list): if not (min(reorder) > 0 and not max(reorder) > len(exprs)): raise ValueError("invalid reorder spec: " + repr(reorder)) js = tuple(j - 1 for j in reorder) s = js + diff(s, js) # 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) # ensure all symbols are filled out rs = set(symbols).difference(*xs).difference(*vs) if rs: exprs.append((None, None, 0)) xs.append(set(rs)) vs.append(set()) ts.append(join(rs, sort=1, enc="<>")) if verbose & self.vI: # output solver information printf("[base={base}, digits={digits}, symbols={symbols!r}, distinct={distinct}]") 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, Mapping): d = dict() for ss in distinct: if sane > 0: if len(set(ss).difference(s2d.keys())) > len(digits): if verbose & self.vW: printf("[SubstitutedExpression: WARNING: distinct=\"{ss}\" has more symbols than available digits]") self._fail = 1 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)) # generate code for the assignment: v = x def assign(v, x, prog, _, in_loop, warn, verbose, self): prog.append(sprintf("{_}try:")) prog.append(sprintf("{_} {v} = {x}")) prog.append(sprintf("{_}except NameError:")) # catch typos in function names prog.append(sprintf("{_} raise")) (extra, msg) = ('', '') if warn and verbose & self.vW: (extra, msg) = (' as e', 'printf("[WARNING: [{x}] {e}]", x=type(e).__name__)') prog.append(sprintf("{_}except Exception{extra}:")) # maybe "except (ArithmeticError, TypeError, ValueError)" if msg: prog.append(sprintf("{_} {msg}")) prog.append(sprintf("{_} {skip}", skip=('continue' if in_loop else 'return'))) # start with any initialisation code if 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(chain(symbols, words), fn=_sym, 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 = ('s' in opt) # use sets in code generation? [currently disabled by default] # 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 Python 3 we could use [[ nonlocal ]] instead of passing the symbols around prog.append(sprintf("{_}def {block}({block_args}):")) _ += indent if decl: prog.append(sprintf("{_}{decl}")) in_loop = False # EXPERIMENTAL: do something about: ": = if ('i' in opt) and 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 , that it not already done # we deal with symbols that appear in the most words first xsyms = sorted(diff(xsyms, done), key=(lambda x: sum(x in w for w in words)), reverse=1) for s in xsyms: # allowable digits for s ds = valid[s] in_loop = True check = None 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 use_sets: check = (sprintf(".difference({check})", check=_set(check)) if check else '') prog.append(sprintf("{_}for {s} in {ds}{check}:", s=_sym(s), ds=_set(ds))) _ += indent else: prog.append(sprintf("{_}for {s} in {ds}:", s=_sym(s))) _ += indent if check: 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: encl(_sym(w), '()'))) assign(vx, sprintf("int({x})"), prog, _, in_loop, warn, verbose, self) # check the value if k == 3: # this is a literal (alphametic) word: # if we have already constructed the word ... if all(x in done for x in val): # we can check it directly prog.append(sprintf("{_}if {vx} == {val}:", val=_sym(val))) _ += indent else: # we extract the symbols from the result # it must have a non-negative value prog.append(sprintf("{_}if {vx} >= 0:")) _ += indent # we can set the value of the word now, to avoid recalculation if len(val) > 1: prog.append(sprintf("{_}{val} = {vx}", val=_sym(val))) # now we check that the value does make the word 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 " + _set(check) 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 additional words which can now be made for w in words: if w != val and 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 if _python == 2: prog.append(sprintf("{_}for {vr} in {block}({block_args}): yield {vr}")) else: prog.append(sprintf("{_}yield from {block}({block_args})")) block = None # [denest] work around statically nested block limit in CPython 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 if decl: prog.append(sprintf("{_}{decl}")) # 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(((encl(s, "'") + ": " + _sym(s)) for s in sorted(done)), sep=', ') if answer: # compute the answer r = _replace_words(answer, (lambda w: encl(_sym(w), '()'))) assign(vr, r, prog, _, in_loop, warn, verbose, self) prog.append(sprintf("{_}yield ({{ {d} }}, {vr})")) else: prog.append(sprintf("{_}yield {{ {d} }}")) # compile any output function definition output_fn = None if output and isinstance(output, basestring): output_fn = solver + "_output_fn" prog.append(output_fn + " = " + output) # 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) gs.update(dict(self=self, cls=self.__class__)) 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 'C' to output code before compilation)") printf("(or use the \"denest=1\" option (--denest, -X) to reduce program complexity)") raise eval(code, gs) # extract the complied code self._solver = gs[solver] if output_fn: self.output = gs[output_fn] 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. """ verbose = (self.verbose if verbose is None else self._verbose(verbose)) if not self._prepared: self._prepare(verbose) if self._fail: return solver = self._solver answer = self.answer header = self.header output = self.output if check is None: check = self.check if first is None: first = self.first if self.sane > 0 and verbose & self.vW and self.template and self.base > len(base_digits()): printf("WARNING: base {self.base} specified, exceeds base_digits()") if verbose & self.vH: if output: # call output with a solution of None to print a header catch(output, (self, None, None) if answer else (self, None)) elif header: printf("{header}") n = 0 for s in solver(): if check and (not check(s)): continue if verbose & self.vT: if output: # call: [[ output(self, soln) ]] or [[ output(self, soln, ans) ]] call(output, (self,) + s if answer else (self, s)) else: self.output_solution((s[0] if answer else s)) # return the result yield s n += 1 if first and first == n: break # solve the puzzle, and just return the answers def answers(self, **kw): """ like solve(), but just return the answer for each solution (assuming the 'answer' parameter has been specified). """ if not self.answer: raise ValueError("'answer' parameter must be specified") for (_, ans) in self.solve(**kw): yield ans # output a solution as: "

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) sequence: >>> icount(primes.between(1, 1000)) 168 """ if t is None: if p is None: # we just need a count of the items in if hasattr(seq, '__len__'): return len(seq) # was: [[ d = collections.deque(enumerate(seq, start=1), maxlen=1); return (d[0][0] if d else 0) ]] return sum(1 for _ in seq) else: return sum(1 for x in seq if p(x)) # otherwise t is defined n = 0 for x in seq: if p is None or p(x): n += 1 if n == t: break return n # icount recipes icount_exactly = lambda seq, p=None, n=None: icount(seq, p, n + 1) == n icount_at_least = lambda seq, p=None, n=None: icount(seq, p, n) == n icount_at_most = lambda seq, p=None, n=None: icount(seq, 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, Mapping): # search the keys # (or we could use find() in the values, and return the correspond index in keys) for (k, x) in seq.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""" n = len(seq) i = find(reversed(seq), v) return (-1 if i == -1 else n - 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: (x, rest) = (seq[0], seq[1:]) for xs in itertools.combinations(rest, n - 1): p = (x,) + xs for ps in _partitions(diff(rest, xs), 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) #yield from fn(seq, n) #[Python 3] for z in fn(seq, n): yield z #[Python 2] # permutations and cycles: # TODO: maybe make a Permutation class (subclass tuple) def cycles(qs, ps=None): """ find the cycles of a permutation (from original order ). if is not specified a sorted version of is used. a cycle of the form (x, y, z) means when the permutation is applied the following transitions occur: x -> y, y -> z, z -> x i.e. y is in the position previously occupied by x, etc. >>> list(cycles([2, 6, 5, 4, 3, 1])) [(1, 2, 6), (3, 5), (4,)] """ if ps is None: ps = sorted(qs) # construct a map of the permutations m = dict(zip(ps, qs)) # construct cycles rs = set(ps) # remaining elements while rs: cyc = [min(rs)] while True: x = m[cyc[-1]] if x != cyc[0]: cyc.append(x) else: break yield tuple(cyc) rs.difference_update(cyc) def permutation(cycs, ps=None): """ permute a sequence according to cycles . if is not specified it is determined from the elements in the cycles. """ if ps is None: ps = sorted(flatten(cycs)) m = dict() for cyc in cycs: m.update(tuples(cyc, 2, circular=1)) return tuple(m[x] for x in ps) # make cycles from elements of length def _cycles(xs, k, ss=[]): if not xs: yield ss else: (x, rest) = (xs[0], xs[1:]) for rs in itertools.permutations(rest, k - 1): #yield from _cycles(diff(rest, rs), k, ss + [(x,) + rs]) #[Python 3] for z in _cycles(diff(rest, rs), k, ss + [(x,) + rs]): yield z #[Python 2] # make permutations from given cycle lengths def _perms_from_cycles(seq, cycs, ss=[]): if not cycs: yield ss else: k = cycs.peek() v = cycs[k] # choose n elements for this group of cycles for xs in itertools.combinations(seq, k * v): for cs in _cycles(xs, k): #yield from _perms_from_cycles(diff(seq, xs), cycs.copy().remove(k, v), ss + cs) #[Python 3] for z in _perms_from_cycles(diff(seq, xs), cycs.copy().remove(k, v), ss + cs): yield z #[Python 2] # generate permutations from specified cycle lengths def permutations_from_cycles(seq, cycs, fn=None): """ generate permutations of sequence using the cycle lengths specified in . if is None then permutations are returned as cycles if you want the permuted sequence use: fn=permutation """ seq = list(seq) # make sure cycs is a multiset if not isinstance(cycs, multiset): cycs = multiset(cycs) assert len(seq) == cycs.sum() # generate permutations as cycles for x in _perms_from_cycles(seq, cycs): yield (fn(x) if fn else x) # see: [ https://enigmaticcode.wordpress.com/2017/05/17/tantalizer-482-lapses-from-grace/#comment-7169 ] def choose(vs, fns, k=None, ss=None, distinct=0, increasing=0, gap=0, multi_vs=0, multi_fns=1): """ 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 in the sequence should be distinct. set 'increasing' if all values in the sequence should be increasing. the 'gap' value can be set to indicate the minimum permissible gap in increasing sequences. if you want to specify a sequence of values for each choice then set the 'multi_vs' flag, and if you want to specify a single function for each test then clear the 'multi_fns' flag. if at least one of the 'multi' flags is set, then the number of choices to perform is set to the length of the corresponding value (if they are both set then the minimum value is used). otherwise the number of choices should be specified as . >>> 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 k is None: if multi_vs: k = len(vs) if multi_fns: k = (len(fns) if k is None else min(len(fns), k)) if k is None: raise ValueError("choose: one of multi_vs, multi_fns, k should be set") if ss is None: ss = list() # are we done? if k == 0: yield ss elif k > 0: i = len(ss) # choose the next value fn = (fns[i] if multi_fns else fns) for v in (vs[i] if multi_vs else vs): if distinct and v in ss: continue if increasing and ss and v < ss[-1] + gap: continue ss_ = ss + [v] if fn is None or fn(*ss_): # choose the rest #yield from choose(vs, fns, k - 1, ss_, distinct, increasing, gap, multi_vs, multi_fns) #[Python 3] for z in choose(vs, fns, k - 1, ss_, distinct, increasing, gap, multi_vs, multi_fns): yield z #[Python 2] def first(seq, count=1, skip=0, fn=list): """ return the first items in sequence (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(primes, 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(powers(0, inf, 2), count=lt(200)) [0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196] """ # anything to skip? if callable(skip): # skip while is true seq = itertools.dropwhile(skip, seq) elif skip > 0: # skip elements [assume non-negative integer] seq = itertools.islice(seq, skip, None) if count == inf or count is None: # keep going to the end of the iterable r = seq elif callable(count): # collect until is false r = itertools.takewhile(count, seq) else: # collect elements [assume positive integer] r = itertools.islice(seq, count) return (r if fn is None else fn(r)) # like first() but default 'fn' is iter ifirst = partial(first, fn=iter) # return the single value if s contains only a single value (else None) # NOTE: similar to the Python expression : [[ [x] = s ]] def singleton(seq, 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(seq, 2, skip) return (r[0] if len(r) == 1 else default) # check container contains only a single value is_singleton = lambda s, skip=0: len(first(s, 2, skip)) == 1 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(inc(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) # repeatedly apply a function to until consecutive values of # measure satisfy function def converge(fn, v, m=identity, cmp=is_equal): """ repeatedly apply function to value until consecutive values of measure satisfy function . For example, finding the square root of 2 using Newton's method: (it would also be appropriate to use: cmp=math.isclose in Python >= 3.5) >>> converge((lambda x, n=2.0: 0.5 * (x + n / x)), 1.0) 1.414213562373095 """ m0 = m(v) while True: v = fn(v) m1 = m(v) if cmp(m0, m1): return v m0 = m1 def uniq(seq, fn=None, seen=(), 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(seen) 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' """ # originally: # for vs in clump(seq, fn=fn): yield vs[0] # this is equivalent, but returns values sooner seq = iter(seq) try: # get the first value v = next(seq) fnv = (v if fn is None else fn(v)) while 1: # return the value yield v # skip repeats while 1: v = next(seq) x = (v if fn is None else fn(v)) if x == fnv: continue # move on to the next value fnv = x break except StopIteration: pass # root: calculate the (positive) nth root of a (positive) real number # we use math.pow rather than **/pow() to avoid generating complex numbers root = lambda x, n: (None if x is None else math.pow(x, n**-1)) # use math.cbrt() if available [Python 3.11+] @static(impl=getattr(math, 'cbrt', None)) def cbrt(x): """ return the cube root of a number (as a float). see also: math.cbrt() (Python 3.11+) >>> round(cbrt(27.0), 6) 3.0 >>> round(cbrt(-27.0), 6) -3.0 """ if x is None: return None fn = cbrt.impl r = (fn(abs(x)) if fn else 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 # generate prime factors # basis = [2, 3, 5] #_prime_factors_deltas = [[2, 1, 2, 2], [4, 2, 4, 2, 4, 6, 2, 6]] #_prime_factors_js = [1, 1] _prime_factors_deltas0 = [2, 1, 2, 2] _prime_factors_deltas1 = [4, 2, 4, 2, 4, 6, 2, 6] def prime_factors(n, limit=inf, validate=0): """ generate prime_factors of . factors are generated in order, and repeated factors will be returned the appropriate number of times. >>> list(prime_factors(60)) [2, 2, 3, 5] >>> list(prime_factors(factorial(12))) [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 5, 5, 7, 11] >>> list(prime_factors(factorial(12) + 1)) [13, 13, 2834329] """ if validate: n = as_int(n, include="0+") if n < 2: return # start with the first list of deltas deltas = _prime_factors_deltas0 x = done = 0 while True: for i in deltas: x += i if x > limit: return while True: if x * x > n: # anything remaining is prime if not (n == 1 or n > limit): yield n return (d, r) = divmod(n, x) if r > 0: break yield x n = d deltas = _prime_factors_deltas1 # for large numbers with large prime factors use prime_factor_h() or sympy.ntheory.factorint() def prime_factor(n, limit=inf, validate=0): """ generate (, ) pairs in the prime factorisation of positive integer . no pairs are returned for 1 (or for non-positive integers). if 'limit' is specified no prime factors greater than this will be returned. for numbers with large prime factors it will take a long time to find them. in this case you probably want to use prime_factor_h() instead. >>> 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)] >>> list(prime_factor(factorial(12) + 1)) [(13, 2), (2834329, 1)] """ for vs in clump(prime_factors(n, limit=limit, validate=validate)): yield (vs[0], len(vs)) # maybe should be called factors() or factorise() def factor(n, fn=prime_factor, validate=0): """ 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()). see also: prime_factors() >>> factor(101) [101] >>> factor(1001) [7, 11, 13] >>> factor(12) [2, 2, 3] >>> factor(125) [5, 5, 5] """ if validate: n = as_int(n, include="0+") 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, validate=0): """ return the divisors of positive integer pow(, ) as a sorted list. the function used to find prime factors can be specified as . if you are factorising large numbers with large prime factors, then you will probably want to provide a function based on prime_factor_h(). >>> divisors(36) [1, 2, 3, 4, 6, 9, 12, 18, 36] >>> divisors(101) [1, 101] >>> ds = divisors(factorial(23) - 1, fn=(lambda n: prime_factor_h(n, mr=1))) >>> print(seq2str(ds)) (1, 51871, 498390560021687969, 25852016738884976639999) """ if validate: n = as_int(n, include="0+") if n == 0 and k > 0: return [0] fs = (fn(n) if callable(fn) else fn) return multiples(fs, k=k) def divisors_pairs(n, k=1, fn=prime_factor, every=0, validate=0): """ generate divisors pairs (a, b) with a <= b, such that a * b = 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 validate: (n, k) = (as_int(n, include="0+"), as_int(k, include="0+")) 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=()): if k == 1: if not (s and n < s[-1]): yield s + (n,) elif k > 1: for (a, b) in divisors_pairs(n): if not (s and a < s[-1]): #yield from _divisors_tuples(b, k - 1, s + (a,)) #[Python 3] for z in _divisors_tuples(b, k - 1, s + (a,)): yield z #[Python 2] def divisors_tuples(n, k, validate=0): """ 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 validate: (n, k) = (as_int(n, include="0+"), as_int(k, include="0+")) if k == 0: if n == 1: return iter([()]) return iter([]) elif k == 1: return iter([(n,)]) elif k == 2: return divisors_pairs(n) elif k > 2: return _divisors_tuples(n, k) def _factorisations(n, ds, i, ss=[]): # are we done? if n == 1: yield tuple(ss) else: # look for the next divisor for j in irange(i, 0, step=-1): d = ds[j] if n >= d: (r, x) = divmod(n, d) if x == 0: #yield from _factorisations(r, ds, j, [d] + ss) #[Python 3] for z in _factorisations(r, ds, j, [d] + ss): yield z #[Python 2] def factorisations(n, fn=divisors, validate=0): """ generate factorisations of (non-negative integer) (sometimes known as: multiplicative partitions) the function to determine the divisors of can be specified as . if is true, then will be validated as a non-negative integer. >>> sorted(factorisations(24), reverse=1) [(24,), (4, 6), (3, 8), (2, 12), (2, 3, 4), (2, 2, 6), (2, 2, 2, 3)] >>> sorted(factorisations(factorial(9) + 1), reverse=1) [(362881,), (269, 1349), (71, 5111), (19, 19099), (19, 71, 269)] """ if validate: n = as_int(n, include="0+") # always return the trivial factorisation yield (n,) if n < 4: return # look for other factorisations using divisors (other than 1 and n) ds = trim(fn(n), head=1, tail=1) #yield from _factorisations(n, ds, len(ds) - 1) #[Python 3] for z in _factorisations(n, ds, len(ds) - 1): yield z #[Python 2] # see also: is_prime_mr(), Primes.is_prime(), gmpy2.is_prime() def is_prime(n, validate=0): # type: (int, int) -> bool """ return True if the non-negative integer is prime. if the argument will be validated as a non-negative integer. 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 is None: return None if validate: n = as_int(n, include="0+") if n < 6: return (n == 2 or n == 3 or n == 5) # 2, 3, 5 -> T; 0, 1, 4 -> F r = n % 6 if r != 1 and r != 5: return False # (n % 6) != (1, 5) -> F # faster to just check divisors % 6 = (1, 5) if n % 5 == 0: return False # was: [[ return n < 49 or all(n % x != 0 and n % (x + 4) != 0 for x in irange(7, isqrt(n), step=6)) ]] if n > 48: for x in irange(7, isqrt(n), step=6): if n % x == 0 or n % (x + 4) == 0: return False return True ## used to be: #for p in prime_factors(n): # return p == n #return False prime = is_prime def is_composite(n, validate=0): "check the non-negative integer is the product of more than 2 primes" if n is None: return None if validate: n = as_int(n, include="0+") if n < 4: return False return not is_prime(n) # Miller-Rabin primality test (originally suggested by Brian Gladman) randrange = lazy_import('random.randrange') def _mr_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, validate=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(1000000000000066600000000000001) 1 >>> is_prime_mr(332306998946228968225951765070086171) 0 """ if n is None: return None if validate: n = as_int(n, include="0+") # 2, 3 = definitely prime; 0, 1, 4 = not prime if n < 5: return (2 if n == 2 or n == 3 else 0) # all other primes have a residue mod 6 of 1 or 5 r = n % 6 if r != 1 and r != 5: return 0 # not prime # 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 _mr_is_composite(9345883071009581737 % n, d, n, s) else 2) # 2 bases are completely accurate for n < 1050535501 if n < 1050535501: return (0 if _mr_is_composite(336781006125 % n, d, n, s) or _mr_is_composite(9639812373923155 % n, d, n, s) else 2) # test remaining numbers with the 7 base set for a in (2, 325, 9375, 28178, 450775, 9780504, 1795265022): # definitely composite if _mr_is_composite(a % n, d, n, s): 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: for _ in range(r): # definitely composite if _mr_is_composite(randrange(2, n - 1), d, n, s): 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)] """ if ps is None and mr == 0: raise ValueError("no heuristics specified") # 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.contains(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, validate=0): """ count the number of divisors of a positive integer . tau(n) = len(divisors(n)) # (but faster) >>> tau(factorial(12)) 792 >>> tau(factorial(23) - 1, fn=(lambda n: prime_factor_h(n, mr=1))) 4 """ if validate: n = as_int(n, include="0+") fs = (fn(n) if callable(fn) else fn) return multiply(e + 1 for (_, e) in fs) 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 """ fs = (fn(n) if callable(fn) else fn) return n > 0 and all(e == 1 for (_, e) in fs) 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 fs = (fn(n) if callable(fn) else fn) for (p, e) in fs: 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 ((b + b - a, b), (a + a + b, a), (b + b + a, b)): if fn(p): _push(ps, p) # Pythagorean Triples: # formerly used the Barning-Hall tree # [ https://en.wikipedia.org/wiki/Tree_of_primitive_Pythagorean_triples ] # now uses Euclid's formula (suggested by Frits) # [ 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 inf (or None), then triples will be generated indefinitely # if order is true, then triples will be returned in order def _pythagorean_primitive(Z=inf, order=0): if order: # use a heap to merge steams of per-m values (increasing) from Euclid's formula from heapq import (heapify, heappush, heappop) ts = list() heapify(ts) # push a (c, b, a, m, n) tuple from (m, n), with c <= Z def push(ts, m, n, Z): (m2, n2) = (m * m, n * n) c = m2 + n2 if c > Z: return (a, b) = (m2 - n2, 2 * m * n) heappush(ts, ((c, a, b, m, n) if b < a else (c, b, a, m, n))) (m0, mv) = (2, 5) # mv = sq(m0) + 1 while True: # bring in any stream that could beat the current top while (not ts) or (mv <= ts[0][0]): if mv > Z: break n0 = 1 + m0 % 2 push(ts, m0, n0, Z) m0 += 1 mv += 2 * m0 - 1 if not ts: return # process the top of the heap (c, b, a, m, n) = heappop(ts) yield (a, b, c) # advance this m-stream for n1 in irange(n + 2, m - 1, step=2): if gcd(m, n1) == 1: push(ts, m, n1, Z) break else: # if we don't care about order we can use a faster formulation of Euclid's formula for m in irange(2, (inf if Z is inf else isqrt(Z))): m2 = m * m # "... not both odd..." for n in irange(1 + m % 2, m - 1, step=2): # "... with m and n coprime ..." if gcd(m, n) == 1: n2 = n * n c = m2 + n2 if c > Z: break a = m2 - n2 b = 2 * m * n yield ((b, a, c) if b < a else (a, b, c)) # generate pythagorean triples (x, y, z) with hypotenuse not exceeding Z def _pythagorean_all(Z, order=0): if Z is inf: order = 1 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): (c, b, a, c0, b0, a0) = heappop(ms) yield (a, b, c) c1 = c + c0 if c1 <= Z: heappush(ms, (c1, b + b0, a + a0, c0, b0, a0)) # return (x, y, z) yield (x, y, z) # add in the next multiple of (x, y, z) z1 = z + z if z1 <= Z: heappush(ms, (z1, y + y, x + x, z, y, x)) # return any remaining multiples while ms: (c, b, a, c0, b0, a0) = heappop(ms) yield (a, b, c) c1 = c + c0 if c1 <= Z: heappush(ms, (c1, b + b0, a + a0, c0, b0, a0)) 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) # return a complete sorted finite list (faster than generating in order) def pythagorean_triples_sorted(n, primitive=0): if n is None: n = inf assert (n < inf) ts = list(pythagorean_triples(n, primitive=primitive, order=0)) ts.sort(key=lambda t: t[::-1]) return ts # 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 @static(sorted=pythagorean_triples_sorted) def pythagorean_triples(n=inf, 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) n can be inf, and triples will be generated indefinitely (although it will eventually exhaust 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 n is None: n = inf # do we want only primitive triples? if primitive: return _pythagorean_primitive(n, order=order) # otherwise also include non-primitive return _pythagorean_all(n, order=order) def fib(*s, **kw): """ generate Fibonacci type sequences (or other recurrence relations) The initial terms are provided as sequence , and subsequent terms are calculated as a function of the preceding 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 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.pop('fn', sum) if kw: raise TypeError(str.format("fib: unknown arguments {kw}", kw=seq2str(kw.keys()))) 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 = intr(n**(1.0 / m)) # return (i if i**m == n else None) # but here we use Newton's method, 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) """ if n is None or n < 0 or k < 1: return None if n == inf: return inf if n >> k == 0: return (0 if n < 1 else 1) n = int(n) # just the integer part # use Newton's method x = 1 << divc(n.bit_length(), k) # initial guess k1 = k - 1 while True: y = (k1 * x + n // (x**k1)) // k if y >= x: return x x = y def irootrem(n, k): r = iroot(n, k) return (r, n - r ** k) 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(16807, 5) 7 >>> is_power(16807, 7) is not None False >>> is_power(0, 7) 0 >>> n = (2**60 + 1) >>> (is_power(n**5, 5) is not None, is_power(n**5 + 1, 5) is not None) (True, False) >>> (is_power(n**7, 7) is not None, is_power(n**7 + 1, 7) is not None) (True, False) """ # remove factors of 2 and 3 while True: if n is None or k == 1: return n (k_, r) = divmod(k, 2) if r == 0: (n, k) = (is_square(n), k_) continue (k_, r) = divmod(k, 3) if r == 0: (n, k) = (is_cube(n), k_) continue break # use iroot() for what is left r = iroot(n, k) if r is None or r**k != n: return None return r is_power_p = (lambda n, k: is_power(n, k) is not None) def sqrt(a, b=None): """ return the (real) square root of / (or just if is None) as a float. >>> sqrt(9) 3.0 >>> sqrt(9, 4) 1.5 """ if a is None: return None # / is operator.truediv() here return math.sqrt(a if b is None else (a / b)) # sq = lambda x: x * x def sq(x, y=None, div=fdiv): "return (x/y) squared (or x squared if y is None)" if y is not None: x = div(x, y) return x * x def diffsq(x, y): "diffsq(x, y) = x^2 - y^2"; return (x - y) * (x + y) 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 non-negative integer n. See also: math.isqrt (Python 3.8+), gmpy2.isqrt(). >>> isqrt(9) 3 >>> isqrt(15) 3 >>> isqrt(16) 4 >>> isqrt(17) 4 """ if n is None or n < 0: return None if n == inf: return inf if n < 4: return (0 if n < 1 else 1) n = int(n) # just the integer part 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) # returns (isqrt(n), remainder) def isqrtrem(n): r = isqrt(n) return (r, n - r * r) # square root floor and ceiling functions sqrtf = isqrt def sqrtc(n): "calculate intc(sqrt(n)), for integer n."; return (isqrt(n) if n < 1 else isqrt(n - 1) + 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, validate=0): # type: (int | NoneType, 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. if is set, then the input value will be validated as a non-negative integer (and an ValueError thrown if it isn't), otherwise the input is assumed to be an integer value, or 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 >>> (is_square(5**46) is not None, is_square(5**46 + 1) is None) (True, True) """ if validate: n = as_int(n, include="0+") # validate -> check for non-negative integer if n is None or n < 0: return None if n < 2: return (n if n == 0 or n == 1 else None) # 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 n is None: return 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) # for k=2 and large n can also use [[ pells.diop_quad(1, 1, n) ]] 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 = sq(a) + sq(b) + ... min_v - specifies the minimum allowable value in the returned sequences sep - specified the minimum separation between values note: for large n and k=2 you can use pells.diop_quad(1, 1, n). >>> 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: max_i2 = ((n - min_v * min_v) if min_v > 0 else n) (i, j) = (isqrt(max_i2), min_v) if i is None: return while not (i < j): r = compare(i * i + j * j, n) if r == 0 and not (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 #yield from sum_of_squares(n - x2, k - 1, x + sep, sep, ss + [x]) #[Python 3] for z in sum_of_squares(n - x2, k - 1, x + sep, sep, ss + [x]): yield z #[Python 2] # generate powers from a range def powers(a, b, k=2, step=1, fn=None): """ generate powers 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)) # generate integers that are perfect powers (n = pow(x, y), x >= 0, y >= 2) def ipowers(exps=None): """ generate, in increasing order, without repeats, non-negative integers that are perfect powers (i.e. n = pow(x, y), x >= 0, y >= 2). is a generator used for exponents (starting from 3). >>> first(ipowers(), 14) [0, 1, 4, 8, 9, 16, 25, 27, 32, 36, 49, 64, 81, 100] >>> first(ipowers(), skip=lt(10), count=lt(100)) # 2-digit powers [16, 25, 27, 32, 36, 49, 64, 81] >>> sum(first(ipowers(), 1000)) 260908296 """ # powers less than 4 yield 0 yield 1 # powers from 4 upwards from heapq import (heapify, heappush, heappop) hi = 1 maxe = 2 pows = [(4, 2, 2)] heapify(pows) # default is to use primes > 2, but irange(3, inf, step=2) also works if exps is None: exps = enigma.primes.generate(3) while True: # find the next power (p, b, e) = heappop(pows) if p > hi: yield p hi = p heappush(pows, ((b + 1)**e, b + 1, e)) # do we need to add in a new exponent? if b == 2: max_e = next(exps) heappush(pows, (2**max_e, 2, max_e)) # search exponents to find perfect powers (this avoids factorisation) def is_ipower(n, max_e=inf, exps=None, validate=0): """ check non-negative integer is a (non-trivial) perfect power. can specify the maximum exponent considered. can provide a sequence of exponents to consider. Note: if you have the decomposition of into prime factors (perhaps from the result of calling prime_factor(n)), say in the list then the following expression will tell you if is a perfect power: call(mgcd, (e for (_, e) in fs)) > 1 >>> is_ipower(64) True >>> is_ipower(63) False """ if validate: n = as_int(n, include="0+") if n is None or n < 0: return None if n < 2: return True if exps is None: exps = enigma.primes.generate(2) max_e = min(n.bit_length() - 1, max_e) # consider increasing exponents for e in exps: if e > max_e: break if is_power(n, e) is not None: return True return False # compose functions in order (forward functional composition, "and then") # so: fcompose(f, g, h)(x) == h(g(f(x))) # [maybe implement this as a class?] 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] fn = lambda *args, **kw: g(f(*args, **kw)) else: # general case def fn(*args, **kw): r = f(*args, **kw) for g in gs: r = g(r) return r fn.fns = (f,) + gs return fn # alias for "function pipelines" pipeline = fcompose # 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))) # see: operator.is_none, operator.is_not_none (Python 3.14+) is_not = operator.not_ is_none = (lambda x: x is None) 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) icbrt = partial(iroot, k=3) def icbrtrem(n): r = iroot(n, 3) return (r, n - r * r * r) # 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, validate=0): """ 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 None True >>> is_cube(0) 0 """ if validate: n = as_int(n, include="0+") if n < 0: return None if n < 2: return n if not is_cube.residues: is_cube.residues = set((i**3) % 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 = iroot(n, 3) if z is not None and z**3 != n: z = None 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, validate=0): """ 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, validate=validate) return (None if r is None else -r) return is_cube(n, validate=validate) # keep the old names as aliases power = is_power cube = is_cube square = is_square # trigonometric functions that work in "half-turns", currently these are approximations # but may be replaced by more accurate versions in the future (Python 3.16) half_turn = 1 # = 1 for angles in half-turns # = 180 for angles in degrees # = pi for angles in radians # convert an angle from "half-turns" to radians def ht2rad(x, ht=None): if ht is None: ht = half_turn if ht != pi: x *= pi if ht != 1: x /= float(ht) return x def rad2ht(x, ht=None): if ht is None: ht = half_turn if ht != pi: if ht != 1: x *= ht x /= pi return x def ht2deg(x, ht=None): if ht is None: ht = half_turn if ht != 180: x *= 180 if ht != 1: x /= float(ht) return x def deg2ht(x, ht=None): if ht is None: ht = half_turn if ht != 180: if ht != 1: x *= ht x /= 180.0 return x # provide trig functions def sin_ht(x, ht=None): return math.sin(ht2rad(x, ht=ht)) def cos_ht(x, ht=None): return math.cos(ht2rad(x, ht=ht)) def tan_ht(x, ht=None): return math.tan(ht2rad(x, ht=ht)) def asin_ht(x, ht=None): return rad2ht(math.asin(x), ht=ht) def acos_ht(x, ht=None): return rad2ht(math.acos(x), ht=ht) def atan2_ht(y, x, ht=None): return rad2ht(math.atan2(y, x), ht=ht) def drop_factors(n, k, validate=0): """ remove factors of from . return (i, m) where n = (m)(k^i) such that m is not divisible by k. both and should be non-negative integers. """ if validate: (n, k) = (as_int(n), as_int(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, validate=0): """ check is a power of . returns such that pow(k, m) = n or None. both and should be non-negative integers. >>> 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 validate: (n, k) = (as_int(n, include="0+"), as_int(k, include="0+")) if n == 0: return (1 if k == 0 else None) if n == 1: return 0 if n < 0 or k < 2: return None (i, m) = drop_factors(n, k) return (i if m == 1 else None) # alias (but see also ndigits()) def ilog(x, base=0, validate=0): """ return the largest positive integer such that: pow(, ) <= . """ if base == 0: base == radix if validate: base = as_int(base, include="+") if x == inf or x is None: return x if x < 1: raise ValueError("ilog: value must be 1+") return ndigits(int(x), base=base) - 1 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 >>> round(trirt(49), 8) 9.4121138 """ 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 tri(k) == n. if is not a triangular number, returns None. >>> is_triangular(5050) 100 >>> is_triangular(66066) 363 >>> 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) is_triangular_p = (lambda x: is_triangular(x) is not None) def digrt(n, base=0, validate=0): """ return the (additive) digital root of positive integer . equivalent to summing the digits of the number, and then repeating the process on the result until a single digit number is reached. >>> digrt(123456789) 9 >>> digrt(sum([1, 2, 3, 4, 5, 6, 7, 8, 9])) 9 >>> digrt(factorial(100)) 9 """ if base == 0: base = radix n = (as_int(n, include="0+") if validate else int(n)) return (0 if n == 0 else int(1 + (n - 1) % (base - 1))) def repdigit(n, d=1, base=0): """ 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 """ if base == 0: base = radix if not (0 <= d < base): raise ValueError(str.format("repdigit: invalid digit: {d!r} for base {base!r}", d=d, base=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. ihypot() is provided to return only integer values. See also: math.hypot() (Python 3.8+). >>> hypot(3, 4) 5.0 >>> hypot(3, 4, 12) 13.0 >>> ihypot(3, 4) 5 """ root = kw.pop('root', math.sqrt) if kw: raise TypeError(str.format("hypot: unknown arguments {kw}", kw=seq2str(kw.keys()))) return root(sum(v * v for v in vs)) # alias for: hypot(..., root=is_square) ihypot = partial(hypot, root=is_square) ihypot.__doc__ = "as hypot(), but return an integer value, or None" # root of combined squares: # for example (for positive x, y, z) # hypot(x, y) = rcs(x, y) # if x^2 + y^2 = z^2, we have: y = rcs(z, -x); x = rcs(z, -y); z = rcs(x, y) def rcs(*vs, **kw): """ calculate the "root of combined squares" of values . for a positive value the square of the value is added, for a negative value the square of the value is subtracted, finally the square root of the total (or None) is returned the function used to determine the square root can be specified via the parameter. ircs() is provided to return only integer values. >>> rcs(3, 4) 5.0 >>> rcs(5, -3) 4.0 >>> ircs(5, -3) 4 """ root = kw.pop('root', math.sqrt) if kw: raise TypeError(str.format("rcs: unknown arguments {kw}", kw=seq2str(kw.keys()))) t = 0 for v in vs: if v < 0: t -= v * v else: t += v * v return catch(root, t) # like rcs() but only return integer square roots ircs = partial(rcs, root=is_square) ircs.__doc__ = "as rcs(), but return an integer value, or None" # return roots of the form n/d in the appropriate domain # with domain = 'C', include = '+' (or '-') means "non-zero" def _roots(domain, include, F, *nds): (pos, neg, zero) = (x in include for x in '+-0') for (n, d) in nds: # calculate the root if domain in "CF": x = F(n / d) elif domain == "Q": x = F(n, d) elif domain == "Z": x = as_int(F(n, d), default=None) if x is None: continue # return the root if domain == "C": if x != 0 or zero: yield x elif x > 0: if pos: yield x elif x < 0: if neg: yield x else: if zero: yield x # find roots of a quadratic equation # domain = "Z" (integer), "Q" (rational), "F" (float), "C" (complex float) def quadratic(a, b, c, v=0, domain="Q", include="+-0", F=None): """ find roots of the quadratic equation: a.x^2 + b.x + c = v 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] """ # select the field if domain in "CF": if F is None: F = (complex if domain == 'C' else float) elif domain in "QZ": if F is None: F = Rational() else: raise ValueError(str.format("quadratic: invalid domain {domain!r}", domain=domain)) # if v specified? if v: c -= v # linear equation? if a == 0: if b == 0: raise ValueError("quadratic: invalid linear equation") return _roots(domain, include, F, (-c, b)) # discriminant D = b * b - 4 * a * c if D < 0 and domain != "C": return _roots(domain, '', None) d = -2 * a if 'v' in _PY_ENIGMA: printf("quadratic: a={a} b={b} c={c} -> ({b} +/- sqrt({D})) / {d}") if domain in "CF": if D == 0: return _roots(domain, include, F, (b, d)) r = pow(F(D), 0.5) # NOTE: gmpy.mpq does not allow pow(mpq, float) return _roots(domain, include, F, (b + r, d), (b - r, d)) elif domain in "QZ": r = is_square_q(D, F=F) if r is not None: if D == 0: return _roots(domain, include, F, (b, d)) return _roots(domain, include, F, (b + r, d), (b - r, d)) return _roots(domain, '', None) # find roots of a cubic equation def cubic(a, b, c, d, v=0, domain='Q', include="+-0", F=None): """ find roots of the cubic equation: a.x^3 + b.x^2 + c.x + d = v in the specified domain: "Z" finds integer solutions "Q" finds rational solutions "F" finds float solutions "C" finds complex solutions >>> sorted(round(x, 6) for x in cubic(1, -6, 11, -6, domain="F")) [1.0, 2.0, 3.0] """ if domain in "CF": if F is None: F = (complex if domain == 'C' else float) elif domain in "QZ": if F is None: F = Rational() else: raise ValueError(str.format("cubic: invalid domain {domain!r}", domain=domain)) # is v specified? if v: d -= v # is quadratic? if a == 0: return quadratic(b, c, d, domain=domain, include=include, F=F) # calculate the determinant D = 18 * a * b * c * d - 4 * b * b * b * d + b * b * c * c - 4 * a * c * c * c - 27 * a * a * d * d if D == 0: # repeated real roots p = b * b - 3 * a * c if p == 0: # triple real root return _roots(domain, include, F, (-b, 3 * a)) else: # single real root, double real root return _roots(domain, include, F, (9 * a * d - b * c, 2 * p), # double root (4 * a * b * c - 9 * a * a * d - b * b * b, a * p), # single root ) # are we looking for exact rational roots? if domain in "QZ": return poly_rational_roots([d, c, b, a], domain=domain, include=include, F=F) # proceed with float/complex approximations to real/complex roots (a2, b2) = (a * a, b * b) p = fdiv(b2 - 3 * a * c, 3 * a2) # often this is -p q = fdiv(2 * b2 * b - 9 * a * b * c + 27 * a2 * d, 27 * a2 * a) if D > 0: # 3 distinct real roots r = sqrt(p, 3) t = fdiv(math.acos(fdiv(-3 * q, 2 * p * r)), 3) z = fdiv(two_pi, 3) rs = ((6 * a * r * math.cos(x) - b, 3 * a) for x in (t, t + z, t - z)) return _roots(domain, include, F, *rs) if D < 0: # 1 real root, 2 complex conjugate roots # find the real root qx = q * -0.5 r = sqrt(q * q * 0.25 - fdiv(p * p * p, 27)) (t1, t2, tx) = (cbrt(qx + r), cbrt(qx - r), fdiv(b, 3 * a)) ts = [t1 + t2] if domain == 'C': # also return the complex roots r34j = sqrt(3, 4) * 1j (e1, e2) = (-0.5 + r34j, -0.5 - r34j) # complex cube roots of 1 ts.extend((e1 * t1 + e2 * t2, e2 * t1 + e1 * t2)) rs = ((t - tx, 1) for t in ts) return _roots(domain, include, F, *rs) 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 (sometimes known as: "round half away from zero"). >>> intr(0.0) 0 >>> intr(2.5) 3 >>> intr(-2.5) -3 >>> list((x, intr(x)) for x in irangef(-3.5, +3.5, step=1.0)) [(-3.5, -4), (-2.5, -3), (-1.5, -2), (-0.5, -1), (0.5, 1), (1.5, 2), (2.5, 3), (3.5, 4)] """ neg = 0 if x < 0: (neg, x) = (1, -x) (d, r) = divmod(x, 1) d = int(d) if r >= 0.5: d += 1 return (-d if neg else d) def snapf(f, **kw): """ snap float to an integer, with tolerance (default: 1e-6) if the value cannot be snapped to an integer the value of the parameter is returned, or a ValueError if not has been set. """ t = kw.get('t', 1e-6) i = intr(f) if abs(f - i) <= t: return i if 'default' in kw: return kw['default'] raise ValueError(str.format("snapf: value out of range {f!r} (t={t!r})", f=f, t=t)) 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 . """ if m == 1: return intf(x) # [[ m is 1 ]] gives a SyntaxWarning return x - (x % m) # was: [[ 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 """ neg = 0 if b < 0: (a, b) = (-a, -b) if a < 0: (neg, a) = (1, -a) r = int((a + a + b) // (b + b)) return (-r if neg else r) 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 or a is None: 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 distinct_chars(*vs, **kw): """ check to see if arguments (as strings) consist of distinct characters. return the number of distinct characters, or None if there are repeated characters. string conversion is done using 'str' (or the function specified by the 'fn' parameter). >>> [distinct_chars("hello")] [None] >>> distinct_chars("world") 5 >>> distinct_chars(sq(98)) 4 """ fn = kw.pop('fn', str) if kw: raise TypeError(str.format("distinct_chars: unknown arguments {kw}", kw=seq2str(kw.keys()))) s = join(vs, fn=fn) k = len(s) return (k if len(set(s)) == k else None) # in Python 3 we can use [[ def is_duplicate(*vs, fn=str): ]] def is_duplicate(*vs, **kw): """ check to see if arguments (as strings) contain duplicate characters. string conversion is done using 'str' (or the function specified by the 'fn' parameter). >>> is_duplicate("hello") True >>> is_duplicate("world") False >>> is_duplicate(sq(99)) False """ return distinct_chars(*vs, **kw) is None # or using regexps #return True if re.search(r'(.).*\1', str(s)) else False duplicate = is_duplicate # this avoids creating a reversed copy of the sequence # Note: for unicode strings this will be checking the sequence of codepoints # you might want to convert the string to a sequence of graphemes first def is_palindrome(seq): """ 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] """ if not hasattr(seq, '__len__'): seq = list(seq) j = len(seq) if j < 2: return True # this is faster than: [[ return zip_eq(seq, reversed(seq), first=j // 2) ]] i = 0 j -= 1 while i < j: if seq[i] != seq[j]: return False i += 1 j -= 1 return True # is a number palindromic in base def is_npalindrome(n, base=0): """ check if integer is palindromic in base . >>> is_npalindrome(1230321) True >>> is_npalindrome(10**5000 + 1) True >>> is_npalindrome(57005, base=9) True """ if base == 0: base = radix 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) # generate -digit palindromes (in base ) def npalindromes(n, base=0): """generate -digit palindromes in base .""" if base == 0: base = radix fn = identity (k, r) = divmod(n, 2) if r != 0: k += 1 fn = (lambda ds: ds[:-1]) for ds in subsets(irange(base), size=k, select='M'): if n < 2 or ds[0] != 0: ds = fn(ds) + rev(ds) yield nconcat(ds, base=base) # 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 def avg(seq, div=fdiv): """ calculate the arithmetic mean (average) of the values in . for sequences this is equivalent to: sum(seq) / len(seq) the function used for division can be provided with the 'div' parameter. Note: For calculating the mean of floats, statistics.mean() may be more appropriate >>> avg(irange(1, 10)) 5.5 """ if hasattr(seq, '__len__'): (t, n) = (sum(seq), len(seq)) else: t = n = 0 for x in seq: t += x n += 1 return div(t, n) iavg = partial(avg, div=div) iavg.__doc__ = "as avg(), but return an integer value, or None." # vector dot product: vec_dot(xs, ys, strict=0, fnp=multiply, fns=sum) def vec_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: vec_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). see also: math.sumprod() (Python 3.12+) >>> vec_dot((1, 3, -5), (4, -2, -1)) 3 >>> call(vec_dot, [(1, 3, -5)] * 2) 35 >>> call(vec_dot, [(1, 3, -5)] * 3) -97 """ strict = kw.pop('strict', None) fns = kw.pop('fns', sum) fnp = kw.pop('fnp', multiply) if kw: raise TypeError(str.format("vec_dot: unknown arguments {kw}", kw=seq2str(kw.keys()))) z = (zip(*vs, strict=strict) if strict is not None else zip(*vs)) return fns(map(fnp, z)) dot = vec_dot # calculate def vec_add(*vs, **kw): k = kw.pop('k', 1) strict = kw.pop('strict', None) fn = kw.pop('fn', None) if kw: raise TypeError(str.format("vec_add: unknown arguments {kw}", kw=seq2str(kw.keys()))) r = list() if k == 1: for xs in (zip(*vs, strict=strict) if strict is not None else zip(*vs)): r.append(sum(xs)) else: for xs in (zip(*vs, strict=strict) if strict is not None else zip(*vs)): x0 = xs[0] r.append(x0 + k * (sum(xs) - x0)) return (fn(r) if fn else r) # multiple argument versions of basic operations def add(*vs): "add(a, b, c, ...) = a + b + c + ..."; return sum(vs) def mul(*vs): "mul(a, b, c, ...) = a * b * c * ..."; return multiply(vs) def bit_or(*vs): "bit_or(a, b, c, ...) = a | b | c | ..." r = 0 for v in vs: r |= v return r def bit_xor(*vs): "bit_xor(a, b, c, ...) = a ^ b ^ c ^ ..." r = 0 for v in vs: r ^= v return r def bit_and(r, *vs): "bit_and(a, b, c, ...) = a & b & c & ..." for v in vs: r &= v if not r: break return r def bit_rol(n, k=1, w=None): if k == 0: return n if k < 0: return bit_ror(n, -k, w=w) if w is None: w = n.bit_length() r = w - k mask = (1 << r) - 1 return ((n & mask) << k) | (n >> r) def bit_ror(n, k=1, w=None): if k == 0: return n if k < 0: return bit_rol(n, -k, w=w) if w is None: w = n.bit_length() r = w - k mask = (1 << k) - 1 return ((n & mask) << r) | (n >> k) 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 abs(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, y0, y1) = (1, 0, 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) if a > 0 else (-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 # simplified CRT # solve: x = a (mod m); x = b (mod n) def crt1(a, m, b, n): g = gcd(m, n) if (a - b) % g != 0: return (None, None) (x, y, z) = (m // g, n // g, (m * n) // g) (u, v, _) = egcd(x, y) return ((b * u * x + a * v * y) % z, z) @static(rtype=None, fail=None) def crt(vs, ms=None): """ general Chinese Remainder Theorem solve: x mod mi = ri where vs = [(r1, m1), (r2, m2), ...] return: (x, m) where the solution is x + k.m for integers k alternatively you can pass: vs = [r1, r2, ...], ms = [m1, m2, ...] >>> tuple(crt([(2, 3), (3, 5), (2, 7)])) (23, 105) >>> tuple(crt([(65, 95), (76, 97), (49, 99)])) (639985, 912285) """ if crt.rtype is None: crt.rtype = namedtuple('CRT', 'x mod') crt.fail = crt.rtype(None, None) x = mm = None for (i, (r, m)) in enumerate(vs if ms is None else zip(vs, ms)): if i == 0: (x, mm) = (r, m) else: (x, mm) = crt1(r, m, x, mm) if x is None: return crt.fail return (crt.fail if x is None else crt.rtype(x, mm)) # find roots of integer valued polynomial p(x) = v modulo n^k def poly_roots_mod(p, v=0, cache=1): """ find roots of integer valued polynomial p(x) = v mod n^k. the function is passed the polynomial