Source code for sympy.ntheory.multinomial

from __future__ import print_function, division

from collections import defaultdict
from sympy.core.compatibility import range, as_int


[docs]def binomial_coefficients(n): """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 """ n = as_int(n) 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
[docs]def binomial_coefficients_list(n): """ 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 """ n = as_int(n) 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
[docs]def multinomial_coefficients(m, n): r"""Return a dictionary containing pairs ``{(k1,k2,..,km) : C_kn}`` where ``C_kn`` are multinomial coefficients such that ``n=k1+k2+..+km``. Examples ======== >>> 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} Notes ===== 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 """ m = as_int(m) n = as_int(n) if not m: if n: return {} 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
[docs]def multinomial_coefficients_iterator(m, n, _tuple=tuple): """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) """ m = as_int(m) n = as_int(n) 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])