# Python functions for http://oeis.org/A194602 etc.

Feb 28th, 2016
588
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. from sympy import poly, symbols
2. from sympy.ntheory import npartitions
3.
4.
5. def a231599_row(n):
6.     """ T(n,k) is the coefficient of x^k in Product_{i=1..n} (1-x^i) """
7.     if n == 0:
8.         return [1]
9.     x = symbols('x')
10.     p = 1
11.     for i in range(1, n+1):
12.         p *= poly(1-x**i)
13.     p = p.all_coeffs()
14.     return p[::-1]
15.
16.
17. def a026807(n, k):
18.     """ A026807: T(n,k) = number of partitions of n with all addends >= k """
19.     if k > n:
20.         return 0
21.     elif k > n/2:
22.         return 1
23.     else:
24.         vc = a231599_row(k-1)
25.         t = len(vc)
26.         vp_range = range(n-t, n+1)
27.         vp_range = vp_range[::-1]  # reverse
28.         r = 0
29.         for i in range(0, t):
30.             r += vc[i] * npartitions(vp_range[i])
31.         return r
32.
33.
34. def a026794(n, k):
35.     """ A026794: T(n,k) = number of partitions of n with smallest addend k """
36.     if n == k:
37.         return 1
38.     return a026807(n - k, k)
39.
40.
41. def dictprodsum(pdict):
42.     s = 0
43.     for key in pdict:
44.         s += key * pdict[key]
45.     return s
46.
47.
48. def dictunion(a_input, b):
49.     a = a_input.copy()
50.     for key in b.keys():
51.         if key in a:
52.             a[key] += b[key]
53.         else:
54.             a[key] = b[key]
55.     return a
56.
57.
58. def dict_to_list(pdict):
59.     """ Turns a partition dict into a partition list, e.g. {4: 2, 5: 3, 6: 1} to [4, 4, 5, 5, 5, 6]. """
60.     plist = []
61.     for key in pdict:
62.         plist += [key for i in range(0, pdict[key])]
63.     return sorted(plist)
64.
65.
66. def dict_to_valnum(pdict):
67.     """ Turns a partition dict into the corresponding value in A194602, e.g. {4: 1, 6: 1} to 479. """
68.     if pdict == {}:
69.         return 0
70.     plist = dict_to_list(pdict)
71.     pstr = ''
72.     for entry in plist:
73.         pstr += '0' + (entry-1) * '1'
74.     return int(pstr, 2)
75.
76.
77. def valnum_to_dict(pval):
78.     """ Turns a value of A194602 into the corresponding partition dict, e.g. 479 to {4: 1, 6: 1}. """
79.     if pval == 0:
80.         return {}
81.     elif pval < 0:
82.         return "The number entered was negative."
83.     binstr = "{0:b}".format(pval)  # n converted to binary string
84.     binarr = binstr.split('0')
85.     pdict = {}
89.             return "The number entered was not a value of A194602."
92.         else:
94.     return pdict
95.
96.
97. def dict_to_keynum(pdict_input):
98.     """ Turns a partition dict into the corresponding key of A194602, e.g. {4: 1, 6: 1} to 39. """
99.     pdict = pdict_input.copy()
100.     result = 0
102.     while True:
103.         if pdict == {}:
104.             return result
105.         n = dictprodsum(pdict)
108.             result += a026807(n - i, i)  # number of partitions of ``n`` with smallest addend equal to ``i``
111.
112.
113. def keynum_to_rownum(pkey):
114.     """ Returns the sum of the parts of the partition, i.e. the partitions row number.
115.    E.g. for 94 ({3: 1, 5: 2}) the row number is 13."""
116.     if pkey == 0:
117.         return 0
118.     row = 1
119.     while True:
120.         row += 1
121.         pnum = npartitions(row)
122.         if pkey <= pnum - 1:
123.             return row
124.
125.
127.     """ Returns the smallest addend and the number of partitions with smaller smallest addends.
128.    E.g. for  94 (in row 13) the smad is 3 and the n.o.p. with smaller smad is 91. """
129.     if pkey == 0:
130.         return 1
132.     count = 0
133.     while count <= pkey:
138.
139.
140. def keynum_to_dict(pkey):
141.     """ Turns a key of A194602 into the corresponding partition dict, e.g. 39 to {4: 1, 6: 1}. """
142.
143.     if pkey == 0:
144.         return {}
145.
146.     row = keynum_to_rownum(pkey)
147.
149.
150.     # partition to be found is ``pos``-th among partitions of ``row`` with smallest addend ``smad``, counted from 1
151.     pos = pkey - nop_with_smaller_smad + 1
152.
153.     # step upwards until all addends ``smad`` are gone, recursively use this function on that partitions key number
154.     rem = 1
156.         rem += 1
157.     result1 = {smad: rem}  # the ``rem`` addends ``smad`` that were removed become part of the result to be returned
158.     new_row = row - rem*smad  # the row of the new partition with the smallest addends removed
159.     key_count = 0
160.     for smaller in range(1, smad):
161.         key_count += a026794(new_row, smaller)
162.     new_pkey = key_count + pos - 1
163.     result2 = keynum_to_dict(new_pkey)
164.     return dictunion(result1, result2)
165.
166.
167. def keynum_to_valnum(pkey):
168.     """ Turns a key of A194602 into the corresponding value, e.g. 39 to 479. """
169.     return dict_to_valnum(keynum_to_dict(pkey))
170.
171.
172. def valnum_to_keynum(pval):
173.     """ Turns a value of A194602 into the corresponding key, e.g. 479 to 39. """
174.     return dict_to_keynum(valnum_to_dict(pval))