Coverage for /usr/lib/python3/dist-packages/sympy/ntheory/multinomial.py: 46%
87 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1from sympy.utilities.misc import as_int
4def binomial_coefficients(n):
5 """Return a dictionary containing pairs :math:`{(k1,k2) : C_kn}` where
6 :math:`C_kn` are binomial coefficients and :math:`n=k1+k2`.
8 Examples
9 ========
11 >>> from sympy.ntheory import binomial_coefficients
12 >>> binomial_coefficients(9)
13 {(0, 9): 1, (1, 8): 9, (2, 7): 36, (3, 6): 84,
14 (4, 5): 126, (5, 4): 126, (6, 3): 84, (7, 2): 36, (8, 1): 9, (9, 0): 1}
16 See Also
17 ========
19 binomial_coefficients_list, multinomial_coefficients
20 """
21 n = as_int(n)
22 d = {(0, n): 1, (n, 0): 1}
23 a = 1
24 for k in range(1, n//2 + 1):
25 a = (a * (n - k + 1))//k
26 d[k, n - k] = d[n - k, k] = a
27 return d
30def binomial_coefficients_list(n):
31 """ Return a list of binomial coefficients as rows of the Pascal's
32 triangle.
34 Examples
35 ========
37 >>> from sympy.ntheory import binomial_coefficients_list
38 >>> binomial_coefficients_list(9)
39 [1, 9, 36, 84, 126, 126, 84, 36, 9, 1]
41 See Also
42 ========
44 binomial_coefficients, multinomial_coefficients
45 """
46 n = as_int(n)
47 d = [1] * (n + 1)
48 a = 1
49 for k in range(1, n//2 + 1):
50 a = (a * (n - k + 1))//k
51 d[k] = d[n - k] = a
52 return d
55def multinomial_coefficients(m, n):
56 r"""Return a dictionary containing pairs ``{(k1,k2,..,km) : C_kn}``
57 where ``C_kn`` are multinomial coefficients such that
58 ``n=k1+k2+..+km``.
60 Examples
61 ========
63 >>> from sympy.ntheory import multinomial_coefficients
64 >>> multinomial_coefficients(2, 5) # indirect doctest
65 {(0, 5): 1, (1, 4): 5, (2, 3): 10, (3, 2): 10, (4, 1): 5, (5, 0): 1}
67 Notes
68 =====
70 The algorithm is based on the following result:
72 .. math::
73 \binom{n}{k_1, \ldots, k_m} =
74 \frac{k_1 + 1}{n - k_1} \sum_{i=2}^m \binom{n}{k_1 + 1, \ldots, k_i - 1, \ldots}
76 Code contributed to Sage by Yann Laigle-Chapuy, copied with permission
77 of the author.
79 See Also
80 ========
82 binomial_coefficients_list, binomial_coefficients
83 """
84 m = as_int(m)
85 n = as_int(n)
86 if not m:
87 if n:
88 return {}
89 return {(): 1}
90 if m == 2:
91 return binomial_coefficients(n)
92 if m >= 2*n and n > 1:
93 return dict(multinomial_coefficients_iterator(m, n))
94 t = [n] + [0] * (m - 1)
95 r = {tuple(t): 1}
96 if n:
97 j = 0 # j will be the leftmost nonzero position
98 else:
99 j = m
100 # enumerate tuples in co-lex order
101 while j < m - 1:
102 # compute next tuple
103 tj = t[j]
104 if j:
105 t[j] = 0
106 t[0] = tj
107 if tj > 1:
108 t[j + 1] += 1
109 j = 0
110 start = 1
111 v = 0
112 else:
113 j += 1
114 start = j + 1
115 v = r[tuple(t)]
116 t[j] += 1
117 # compute the value
118 # NB: the initialization of v was done above
119 for k in range(start, m):
120 if t[k]:
121 t[k] -= 1
122 v += r[tuple(t)]
123 t[k] += 1
124 t[0] -= 1
125 r[tuple(t)] = (v * tj) // (n - t[0])
126 return r
129def multinomial_coefficients_iterator(m, n, _tuple=tuple):
130 """multinomial coefficient iterator
132 This routine has been optimized for `m` large with respect to `n` by taking
133 advantage of the fact that when the monomial tuples `t` are stripped of
134 zeros, their coefficient is the same as that of the monomial tuples from
135 ``multinomial_coefficients(n, n)``. Therefore, the latter coefficients are
136 precomputed to save memory and time.
138 >>> from sympy.ntheory.multinomial import multinomial_coefficients
139 >>> m53, m33 = multinomial_coefficients(5,3), multinomial_coefficients(3,3)
140 >>> m53[(0,0,0,1,2)] == m53[(0,0,1,0,2)] == m53[(1,0,2,0,0)] == m33[(0,1,2)]
141 True
143 Examples
144 ========
146 >>> from sympy.ntheory.multinomial import multinomial_coefficients_iterator
147 >>> it = multinomial_coefficients_iterator(20,3)
148 >>> next(it)
149 ((3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), 1)
150 """
151 m = as_int(m)
152 n = as_int(n)
153 if m < 2*n or n == 1:
154 mc = multinomial_coefficients(m, n)
155 yield from mc.items()
156 else:
157 mc = multinomial_coefficients(n, n)
158 mc1 = {}
159 for k, v in mc.items():
160 mc1[_tuple(filter(None, k))] = v
161 mc = mc1
163 t = [n] + [0] * (m - 1)
164 t1 = _tuple(t)
165 b = _tuple(filter(None, t1))
166 yield (t1, mc[b])
167 if n:
168 j = 0 # j will be the leftmost nonzero position
169 else:
170 j = m
171 # enumerate tuples in co-lex order
172 while j < m - 1:
173 # compute next tuple
174 tj = t[j]
175 if j:
176 t[j] = 0
177 t[0] = tj
178 if tj > 1:
179 t[j + 1] += 1
180 j = 0
181 else:
182 j += 1
183 t[j] += 1
185 t[0] -= 1
186 t1 = _tuple(t)
187 b = _tuple(filter(None, t1))
188 yield (t1, mc[b])