Simple random walk with transition probabilities derived from out degree

def compute_random_walk(A, walk_length):
    '''Compute random walk based on transition probabilities
    Params:
        A: Adjacency matrix (numpy array)
        walk_length: The length of the walk (int)
    Returns:
        A list of indices passed by the walk
    '''
    passed_node_indices = []

    # compute the transition probability matrix
    r_sum = np.sum(A, axis=1)
    P = A/r_sum.reshape(-1,1)

    # now extract the nodes based on the walk
    node_indices = np.arange(A.shape[0])
    current_pos = np.random.choice(node_indices, size=1)[0]
    passed_node_indices.append(current_pos)
    for i in range(walk_length):
        pdist = P[current_pos].tolist()
        current_pos = np.random.choice(node_indices, size=1, p=pdist)[0]
        passed_node_indices.append(current_pos)
    assert len(passed_node_indices) == walk_length+1
    return passed_node_indices