def a005901(layer : int) -> int:
"""
https://oeis.org/A005901
layer >= 0
"""
if int(layer) < 0:
raise ValueError
if not (int(layer) == float(layer)):
raise ValueError
if layer == 0:
return 1
return 10 * layer ** 2 + 2
print([a005901(x) for x in range(30)])
[1, 12, 42, 92, 162, 252, 362, 492, 642, 812, 1002, 1212, 1442, 1692, 1962, 2252, 2562, 2892, 3242, 3612, 4002, 4412, 4842, 5292, 5762, 6252, 6762, 7292, 7842, 8412]
def accum_a005901(layer):
"""
A005902 cumulative version of A005901
"""
total = 0
for x in range(layer + 1):
total = total + a005901(x)
return total
Upon reading the notes for A005902 we find there's a more direct way to compute the cumulative number of balls i.e. we do not have to sum the balls in each layer out to some target layer. Enter the target layer $L$ directly into this formula:
$$ S = \frac{(2L + 1)(5L^{2} + 5L + 3)}{3} $$import sympy as sp
L = sp.Symbol('L')
expr = (2*L + 1)*(5*L**2 + 5*L + 3)/3
expr
formula = expr.expand()
formula
def short_cut(layer):
if layer == 0:
return 1
return ((2 * layer + 1) * (5 * layer**2 + 5 * layer + 3)) // 3 # see below
short_cut(10)
3871
print([short_cut(x) for x in range(30)])
[1, 13, 55, 147, 309, 561, 923, 1415, 2057, 2869, 3871, 5083, 6525, 8217, 10179, 12431, 14993, 17885, 21127, 24739, 28741, 33153, 37995, 43287, 49049, 55301, 62063, 69355, 77197, 85609]
int(formula.evalf(subs={L:10}))
3871
print([int(expr.evalf(subs={L:x})) for x in range(30)])
[1, 13, 55, 147, 309, 561, 923, 1415, 2057, 2869, 3871, 5083, 6525, 8217, 10179, 12431, 14993, 17885, 21127, 24739, 28741, 33153, 37995, 43287, 49049, 55301, 62063, 69355, 77197, 85609]
6 // 2.1
2.0
6 // 2
3
6 // 5
1