Liệt kê dãy số theo quy luật với Python

Bài toán: Ứng với mỗi số tự nhiên x (0 < x < 10⁶) ta có số tự nhiên f(x) bằng tổng bình phương các chữ số của x.

Ví dụ: Từ số 12 viết ra dãy số có 4 phần tử. Ta có dãy số theo quy luật trên sẽ là: 12 5 25 29. 

Giải thích: Số đầu tiên là 12 giữ nguyên, số thứ 2 là 1² + 2² = 5, số tiếp theo là 5² = 25, số cuối cùng là 2² + 5² = 29

Yêu cầu. Đọc dãy số tự nhiên x và số phần tử n từ file DAYSO.INP, kết quả dãy số viết theo quy luật sẽ ghi vào file DAYSO.OUT.

Ví dụ:

Code tham khảo:

def f(x):
    st = str(x)
    s = 0
    for i in range(0, len(st)):
        s = s + int(st[i])**2
    return s
f_in = open("DAYSO.INP", "r")
f_out = open("DAYSO.OUT", "w")
line = f_in.readline()
a = line.split();
x, n = int(a[0]), int(a[1])
f_out.write(str(x))
for i in range(0, n-1):
    f_out.write(" " + str(f(x)))
    x = f(x)
f_in.close()
f_out.close()