python计算可达矩阵

从别人那里看到的，原代码貌似有些问题，修改了一下

matrix = [[0, 0, 1, 0, 0, 0, 0, 0, 0],
          [0, 0, 1, 1, 0, 0, 0, 0, 0],
          [0, 0, 0, 0, 1, 0, 0, 0, 0],
          [0, 0, 0, 0, 1, 1, 0, 0, 0],
          [0, 0, 0, 0, 0, 0, 1, 0, 0],
          [0, 0, 0, 0, 0, 0, 1, 1, 0],
          [0, 0, 0, 0, 0, 0, 0, 0, 1],
          [0, 0, 0, 0, 0, 0, 0, 0, 1],
          [0, 0, 0, 0, 0, 0, 0, 0, 0]]
          
def get_reach_matrix(matrix):
    '''
    @description: 从邻接矩阵获得可达矩阵
    @matrix {list} list列表形式的矩阵，因此需要np.mat将其转化为矩阵 
    @return: 01形式的可达矩阵
    @usage: 
        matrix = [[0, 0, 1, 0, 0, 0, 0, 0, 0], 
                [0, 0, 1, 1, 0, 0, 0, 0, 0], 
                [0, 0, 0, 0, 1, 0, 0, 0, 0],
                [0, 0, 0, 0, 1, 1, 0, 0, 0],
                [0, 0, 0, 0, 0, 0, 1, 0, 0], 
                [0, 0, 0, 0, 0, 0, 1, 1, 0], 
                [0, 0, 0, 0, 0, 0, 0, 0, 1],
                [0, 0, 0, 0, 0, 0, 0, 0, 1], 
                [0, 0, 0, 0, 0, 0, 0, 0, 0]]

        get_reach_matrix(matrix)
    '''
    import numpy as np
    A = np.mat(matrix)
    I = np.identity(len(A))
    newMat = A+I
    oldMat = newMat
    flag = 0
    step = 1
    while flag == 0:
        oldMat = newMat
        newMat = oldMat*(A+I)
        for i in range(len(newMat)):
            for j in range(len(newMat)):
                if newMat[i, j] >= 1:
                    newMat[i, j] = 1
        step += 1
        print(step)
        if (oldMat == newMat).all():
            flag = 1
            print(newMat, step)
            
get_reach_matrix(matrix)

输出：

2
3
4
5
[[1. 0. 1. 0. 1. 0. 1. 0. 1.]
 [0. 1. 1. 1. 1. 1. 1. 1. 1.]
 [0. 0. 1. 0. 1. 0. 1. 0. 1.]
 [0. 0. 0. 1. 1. 1. 1. 1. 1.]
 [0. 0. 0. 0. 1. 0. 1. 0. 1.]
 [0. 0. 0. 0. 0. 1. 1. 1. 1.]
 [0. 0. 0. 0. 0. 0. 1. 0. 1.]
 [0. 0. 0. 0. 0. 0. 0. 1. 1.]
 [0. 0. 0. 0. 0. 0. 0. 0. 1.]] 5