As professed, this one was pretty easy. Though I had to let it run for around an hour, I was still able to finish it up in a single day. I implemented what I thought was a pretty slow solution then grabbed the full list of S numbers using that. When I put that sequence into OEIS I found https://oeis.org/A038206. From there I learned that the square number must be 0 or 1 mod 9. That lead me to the solution.
def T(n):
total, r, s = 0, 4, 16
while s <= n:
r += 1
s, mod = r ** 2, r % 9
if mod != 0 and mod != 1:
continue
str_s = str(s)
for c in range(1, 1 << len(str_s) - 1):
splits = ((c >> i) & 1 for i in range(len(str_s)-1))
i = summ = 0
for j, split in enumerate(splits):
if not split:
continue
summ += int(str_s[i:j+1])
if summ > r:
break
i = j+1
else:
summ += int(str_s[i:])
if summ == r:
total += s
break
return total
if __name__ == '__main__':
import sys
p = int(sys.argv[1])
print(T(10 ** p))
import pytest
from problem import T
_test_T = (
(10 ** 1, 0),
(10 ** 2, 181),
(10 ** 3, 181),
(10 ** 4, 41333),
(10 ** 5, 184767),
(10 ** 6, 10804656),
(10 ** 7, 30940313),
(10 ** 8, 2818842841),
(10 ** 9, 6222187932),
)
@pytest.mark.parametrize('n,expect', _test_T)
def test_T(n, expect):
assert expect == T(n)