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