Numpyで行列をちょっと太らせる

シンボル列の間の遷移確率を計算するとか何とかで、種数 * 種数の行列が必要になるんだけども、でも事前に種数をカウントするのは嫌で新しいシンボルにであった時に行列を太らせたいな、と思ったのです。最初は正方行列を2つ取って、斜めにくっつけた行列を返す関数があるんじゃないかと探してみたのだけど、見つけられなかったのでc_とr_で貼り付けました。

def enlarge(mat):
    """
    widen matrix
    >>> x = np.array([[1, 2], [3, 4]])
    >>> x
    array([[1, 2],
           [3, 4]])
    >>> enlarge(x)
    array([[ 1.,  2.,  0.],
           [ 3.,  4.,  0.],
           [ 0.,  0.,  1.]])
    """
    N, M = mat.shape
    assert N == M
    result = np.r_[np.c_[mat, np.zeros(N)], np.zeros((1, N + 1))]
    result[N, N] = 1.0
    return result

追記: 中谷さん案(id:n_shuyo)の方が圧倒的に速い。N = 5 で僕の実装が120 us、中谷さんの実装が8.34 us。N = 100 で僕の実装が178 us、中谷さんの実装が27.9 us。

result = numpy.zeros((N+1, N+1)) 
result[0:N,0:N] = mat 
result[N,N] = 1