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