Two-Step Transition Probability

Problem

Row i of a transition matrix gives the one-step distribution from state i, so the probability of moving from i to j in exactly two steps is the (i, j) entry of the squared matrix. Return the sum over k of matrix[i][k] * matrix[k][j]. Return 0.0 for an empty matrix or an out-of-range or malformed index.

Hint

Sum over the possible intermediate state k.

Test cases

case 1

in: [[[0.5,0.5],[0.2,0.8]],0,0]

expected: 0.35

case 2

in: [[[0.5,0.5],[0.2,0.8]],0,1]

expected: 0.65

case 3

in: [[[1,0],[0,1]],1,0]

expected: 0

case 4

in: [[[0.3,0.7],[0.6,0.4]],1,1]

expected: 0.5800000000000001

case 5

in: [[[0.5,0.5],[0.2,0.8]],5,0]

expected: 0

Starter codePython

def two_step_transition_probability(matrix, i, j):
    # Your code here
    pass

5 test cases · run your code against them in the editor