1.
понравился прием с reshape — покажу на примере: пускай мы интересуемся, скажем, представлением чисел в виде суммы квадратов

import numpy as np

def sum_squares(N,k):
    counts1 = np.zeros(N*N+1,dtype=int)
    counts1[0] = 1
    np.add.at(counts1,[n*n for n in range(1,N+1)],2)
    counts = counts1.copy()
    for _ in range(k-1):
        counts = np.convolve(counts,counts1)
    return counts[:N*N+1]

(такой подсчет сумм сверткой — это примерно как возводить производящую функцию квадратов в степени, только без “развлечений” с символическими вычислениями)

тогда можно печатать ответы по сколько-то в строку в духе

print(np.reshape(sum_squares(20,3)[:240],(-1,8)))

(число 8 можно подобрать, желая чтобы расположение нулей выглядело не так хаотично — и возникающую гипотезу дальше легко доказать)


2.
для сумм 4, скажем, квадратов ничего из reshape вроде бы не видно… зато помогает посмотреть на значения для простых чисел — на бумажке это немного утомительно, а компьютерно легко (все числа там получаются кратными 8, так что сразу на 8 разделим):

counts = sum_squares(N:=10,4)
for p in np.nonzero(is_prime(N*N))[0]:
    print(f"{p}: {counts[p]//8}")

какая возникает гипотеза? )

а что будет для p², p³? (спойлер: ответ равен количеству одномерных подпростравнств в 2-, 3-, 4-мерных пространствах над Z/p — это приятно понимать при помощи кватернионов)

(upd)
заодно посчитал количество представлений простых в виде сумму 8 квадратов — и тут легко угадал и неизвестный мне заранее ответ: 16(p³+1) — мб надо подумать, как это доказать