numpyで一次元配列を下三角行列になるように変形したい

たとえば[x1, x2, x3, x4, x5, x6]という配列を与えた時に以下を返すような関数を作りたいです。

[[x1, 0, 0],
[x2,x4, 0],
[x3,x5,x6]]

これを多様なサイズの配列にも適用出来るようにするにはどのように記述すればよいのでしょうか？

行動規範の内容に同意します

回答2件

ベストアンサー

以下でできます。

python
1>>> import numpy as np
2>>> from math import sqrt, ceil
3>>> eps = 1e-10
4>>> def find_size(length):
5...     return ceil(sqrt(1+8*length)/2 - 0.5 - eps)
6...
7>>> N = 21
8>>> v = np.arange(1, N+1)
9>>> print(v)
10[ 1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21]
11>>> size = find_size(len(v))
12>>> tm = np.zeros(size * size, dtype=v.dtype).reshape(size, size)
13>>> position = [0] + list(np.arange(size, 0,  -1).cumsum())
14>>> for i in range(size):
15...     tm[i:,i] = v[position[i]:position[i+1]]
16...
17>>> print(tm)
18[[ 1  0  0  0  0  0]
19 [ 2  7  0  0  0  0]
20 [ 3  8 12  0  0  0]
21 [ 4  9 13 16  0  0]
22 [ 5 10 14 17 19  0]
23 [ 6 11 15 18 20 21]]

投稿2021/04/03 18:19

ppaul

総合スコア24670

解決済になっていますが、より簡単かつ高速な方法があるので書いておきます。

NumPyには多くの便利な関数が実装されており、for文に頼らない方法を模索することは重要です。
np.tri()を用いて三角行列用マスクを作成することで、元のベクトルを一度に割り当てることができます。
numpy.tri — NumPy Manual

python
1def to_triangle(v):
2    size = int(np.sqrt(v.size * 2))
3    out = np.zeros((size, size), v.dtype)
4    out[~np.tri(size, dtype=bool, k=-1)] = v
5    return out.T

実行例：

python
1In [11]: to_triangle(np.arange(1, 7))
2Out[11]:
3array([[1, 0, 0],
4       [2, 4, 0],
5       [3, 5, 6]])
6
7In [12]: to_triangle(np.arange(1, 22))
8Out[12]:
9array([[ 1,  0,  0,  0,  0,  0],
10       [ 2,  7,  0,  0,  0,  0],
11       [ 3,  8, 12,  0,  0,  0],
12       [ 4,  9, 13, 16,  0,  0],
13       [ 5, 10, 14, 17, 19,  0],
14       [ 6, 11, 15, 18, 20, 21]])

速度比較例：

python
1# 長さ約5000のベクトル -> (100, 100)の三角行列
2In [13]: N = np.arange(101).sum()
3
4In [14]: %timeit to_triangle(np.arange(1, N+1))  # 本回答の提案する方法
537.6 µs ± 1.55 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
6
7In [15]: %timeit by_ppaul(np.arange(1, N+1))  # ベストアンサーの方法
8111 µs ± 1.26 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)