Key recursion to calculate n-step transition probabilities in Markov chains

transition_probabilities.py

# define the transition probabilities of the Markov chain
TP = [
    [0.4, 0.6, 0, 0, 0, 0, 0],
    [0.2, 0, 0.8, 0, 0, 0, 0],
    [0, 0.4, 0.6, 0, 0, 0, 0],
    [0, 0, 0.4, 0, 0.3, 0.3, 0],
    [0, 0, 0, 0, 0, 1, 0],
    [0, 0, 0, 0.2, 0, 0, 0.8],
    [0, 0, 0, 0, 1, 0, 0]
]

# calculates the transition probability from state i to state j in n time steps
def r(i, j, n):
    if n == 1:
        return TP[i][j]
    else:
        result = 0
        for k in range(0, len(TP)):
            result = result + r(i, k, n-1) * TP[k][j]
        return result

# an example: what is r_{0, 2}(3)
# that is: the probability that, starting from state 0, after 3 time steps we end up in state 2
print(r(0, 2, 3))