3b1b.py

challenge_accepted.py

N = 2000
R = 5


def matmul(a, b):
    c = [[0] * R for _ in range(R)]
    for i in range(R):
        for j in range(R):
            for k in range(R):
                c[i][j] += a[i][k] * b[k][j]
    return c


def eye():
    m = [[0] * R for _ in range(R)]
    for i in range(R):
        m[i][i] = 1
    return m


def create_mat(with_diagonal):
    m = [[0] * R for _ in range(R)]
    for i in range(R):
        m[i][i] += 1
        m[i][(i - with_diagonal) % R] += 1
    return m


def matpow(a, p):
    b = eye()
    while p > 0:
        if p % 2 == 1:
            b = matmul(b, a)
        p //= 2
        a = matmul(a, a)
    return b


def main():
    M = eye()

    for i in range(1, 5+1):
        r = create_mat(i)
        p = N // 5 + int(N % 5 >= i)
        M = matmul(M, matpow(r, p))

    print(M[0][0])


if __name__ == '__main__':
    main()