# Source code for sympy.ntheory.multinomial

from collections import defaultdict

def binomial_coefficients(n):
[docs] """Return a dictionary containing pairs :math:{(k1,k2) : C_kn} where :math:C_kn are binomial coefficients and :math:n=k1+k2. Examples ======== >>> from sympy.ntheory import binomial_coefficients >>> binomial_coefficients(9) {(0, 9): 1, (1, 8): 9, (2, 7): 36, (3, 6): 84, (4, 5): 126, (5, 4): 126, (6, 3): 84, (7, 2): 36, (8, 1): 9, (9, 0): 1} See Also ======== binomial_coefficients_list, multinomial_coefficients """ d = {(0, n):1, (n, 0):1} a = 1 for k in range(1, n//2+1): a = (a * (n-k+1))//k d[k, n-k] = d[n-k, k] = a return d def binomial_coefficients_list(n):
[docs] """ Return a list of binomial coefficients as rows of the Pascal's triangle. Examples ======== >>> from sympy.ntheory import binomial_coefficients_list >>> binomial_coefficients_list(9) [1, 9, 36, 84, 126, 126, 84, 36, 9, 1] See Also ======== binomial_coefficients, multinomial_coefficients """ d = [1] * (n+1) a = 1 for k in range(1, n//2+1): a = (a * (n-k+1))//k d[k] = d[n-k] = a return d def multinomial_coefficients0(m, n, _tuple=tuple, _zip=zip):
"""Return a dictionary containing pairs {(k1,k2,..,km) : C_kn} where C_kn are multinomial coefficients such that n=k1+k2+..+km. For example: >>> from sympy import multinomial_coefficients >>> multinomial_coefficients(2, 5) # indirect doctest {(0, 5): 1, (1, 4): 5, (2, 3): 10, (3, 2): 10, (4, 1): 5, (5, 0): 1} The algorithm is based on the following result: Consider a polynomial and its n-th exponent:: P(x) = sum_{i=0}^m p_i x^i P(x)^n = sum_{k=0}^{m n} a(n,k) x^k The coefficients a(n,k) can be computed using the J.C.P. Miller Pure Recurrence [see D.E.Knuth, Seminumerical Algorithms, The art of Computer Programming v.2, Addison Wesley, Reading, 1981;]:: a(n,k) = 1/(k p_0) sum_{i=1}^m p_i ((n+1)i-k) a(n,k-i), where a(n,0) = p_0^n. """ if not m: if n: return {} else: return {(): 1} if m == 2: return binomial_coefficients(n) symbols = [(0,)*i + (1,) + (0,)*(m-i-1) for i in range(m)] s0 = symbols[0] p0 = [_tuple(aa - bb for aa, bb in _zip(s, s0)) for s in symbols] r = {_tuple(aa*n for aa in s0): 1} l = [0] * (n*(m - 1) + 1) l[0] = list(r.items()) for k in range(1, n*(m - 1) + 1): d = defaultdict(int) for i in range(1, min(m, k + 1)): nn = (n + 1)*i - k if not nn: continue t = p0[i] for t2, c2 in l[k - i]: tt = _tuple([aa + bb for aa, bb in _zip(t2, t)]) d[tt] += nn*c2 if not d[tt]: del d[tt] r1 = [(t, c//k) for (t, c) in d.items()] l[k] = r1 r.update(r1) return r def multinomial_coefficients(m, n):
[docs] r"""Return a dictionary containing pairs {(k1,k2,..,km) : C_kn} where C_kn are multinomial coefficients such that n=k1+k2+..+km. For example: >>> from sympy.ntheory import multinomial_coefficients >>> multinomial_coefficients(2, 5) # indirect doctest {(0, 5): 1, (1, 4): 5, (2, 3): 10, (3, 2): 10, (4, 1): 5, (5, 0): 1} The algorithm is based on the following result: .. math:: \binom{n}{k_1, \ldots, k_m} = \frac{k_1 + 1}{n - k_1} \sum_{i=2}^m \binom{n}{k_1 + 1, \ldots, k_i - 1, \ldots} Code contributed to Sage by Yann Laigle-Chapuy, copied with permission of the author. See Also ======== binomial_coefficients_list, binomial_coefficients """ if not m: if n: return {} else: return {(): 1} if m == 2: return binomial_coefficients(n) if m >= 2*n and n > 1: return dict(multinomial_coefficients_iterator(m, n)) t = [n] + [0] * (m - 1) r = {tuple(t): 1} if n: j = 0 # j will be the leftmost nonzero position else: j = m # enumerate tuples in co-lex order while j < m - 1: # compute next tuple tj = t[j] if j: t[j] = 0 t[0] = tj if tj > 1: t[j + 1] += 1 j = 0 start = 1 v = 0 else: j += 1 start = j + 1 v = r[tuple(t)] t[j] += 1 # compute the value # NB: the initialization of v was done above for k in range(start, m): if t[k]: t[k] -= 1 v += r[tuple(t)] t[k] += 1 t[0] -= 1 r[tuple(t)] = (v * tj) // (n - t[0]) return r def multinomial_coefficients_iterator(m, n, _tuple=tuple):
[docs] """multinomial coefficient iterator This routine has been optimized for m large with respect to n by taking advantage of the fact that when the monomial tuples t are stripped of zeros, their coefficient is the same as that of the monomial tuples from multinomial_coefficients(n, n). Therefore, the latter coefficients are precomputed to save memory and time. >>> from sympy.ntheory.multinomial import multinomial_coefficients >>> m53, m33 = multinomial_coefficients(5,3), multinomial_coefficients(3,3) >>> m53[(0,0,0,1,2)] == m53[(0,0,1,0,2)] == m53[(1,0,2,0,0)] == m33[(0,1,2)] True Examples ======== >>> from sympy.ntheory.multinomial import multinomial_coefficients_iterator >>> it = multinomial_coefficients_iterator(20,3) >>> next(it) ((3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), 1) """ if m < 2*n or n == 1: mc = multinomial_coefficients(m, n) for k, v in mc.items(): yield(k, v) else: mc = multinomial_coefficients(n, n) mc1 = {} for k, v in mc.items(): mc1[_tuple(filter(None, k))] = v mc = mc1 t = [n] + [0] * (m - 1) t1 = _tuple(t) b = _tuple(filter(None, t1)) yield (t1, mc[b]) if n: j = 0 # j will be the leftmost nonzero position else: j = m # enumerate tuples in co-lex order while j < m - 1: # compute next tuple tj = t[j] if j: t[j] = 0 t[0] = tj if tj > 1: t[j + 1] += 1 j = 0 else: j += 1 t[j] += 1 t[0] -= 1 t1 = _tuple(t) b = _tuple(filter(None, t1)) yield (t1, mc[b])