【Python】 a[x][y]といった構文で、複数インデックスを受け取り、値を返す多次元配列風クラスを自作したい。

Question

前提・実現したいこと

a[x][y]といった構文で、複数インデックスを受け取り、値を返す多次元配列風クラスを自作したいと考えています。

既に存在している数値計算プログラムの中で、N x N x N x Nという4次元正方行列を取扱っています。通常使用では問題なかったのですが、最近Nのサイズが大きくなってしまう計算を取り扱う必要が出てきて、メモリサイズ的に厳しくなってしまいました。

よくよくコードを見直すと、アクセスはそこまで頻繁でなく、行列に収められた値の種類は非常に限られることから、、例えばa[k][l][m][n]というアクセスに対し、k,l,m,nという4つのインデックスを適当に演算して、もっとコンパクトな他の行列に収められた数字を渡す形にすれば、大幅にメモリサイズを圧縮できることがわかりました。

下記にも示すように、＿getitem＿を用いて、タプルで複数インデックスを受け取るのがベストプラクティスのようですが、もし可能であれば、既存コードへの変更を最小限にするためにも、a[k,l,m,n]ではなく、a[k][l][m][n]で実装できればありがたいです。

参考にしたサイト

ググってみた結果、例えば次のような質問と回答をみつけました。
reddit: getitem method in 2D array class.

回答者は、二番目の選択肢として「できる」と言っているようなのですが、お恥ずかしながら、私はうまく噛み砕いて実装できませんでした。こちらについて、追加の解説を頂ければ嬉しいです。

Some clues.

Version one, using a tuple:

def __getitem__(self, row_col_as_tuple):
    row = .... ?
    col = .....? #How would you obtain `row` and `column` from `row_col_as_tuple`?
    return ??

Version two, allowing list-like access eg myarray[1][2]. Think what this is actually doing; it first calls myarray[1] and then calls [2] on whatever myarray[1] returns. That is, it is actually making two calls to getitem, and consequently, whatever object getitem returns, that object itself must be a class that with a getitem method.

def __getitem__(self, key): # where `key` in an integer
    value = self._arr[?]
    return ?? # Here you have to return a class with a __getitem__ method that has access to `self._arr`.

他の参考サイト:
stack overflow: Is there a cleaner way to use a 2 dimensional array?

Answer 1

a[i][j][k][l]を参照した際、何らかの計算を行いたいということでしょうか。
かなり無理やりな方法であれば実現はできます。Wandbox

def func(src):
    return sum(src)  # とりあえず例として和を計算

class Arr:
    def __init__(self, *, indice=None):
        self._indice = indice or []
    
    def __getitem__(self, idx):
        indice = self._indice + [idx]

        if len(indice) == 4:
            return func(indice)
        
        return Arr(indice=indice)
    
arr = Arr()
print(arr[1][2][3][4])   # => 10 と出力される

問題点：

• 逐一オブジェクトを作るオーバーヘッド
• 同じクラスのオブジェクトなのに、getitemが異なる型のオブジェクトを返すのは分かり辛い

---

断言はできないですけれど、[i][j][k][l] を [i,j,k,l] に変えた方が後々楽な気がします。
エディタで置換したら案外一発で動作するかもしれないですし、試す価値はあります。

Comments

[goldberg]

ありがとうございます。このような再帰的な書き方ができるとは目からウロコでした。
ちなみに実装して速度を比較したところ、もちろんfunc(src)内部の演算の不可次第ですが、私の環境では[i,j,k,l]とインデックスを取るのに比べて倍程度の時間がかかるようです。これをどう捉えるかは微妙なところですが。

Answer 2

こういうことらしいです。

class Sparse2DArray:
    class ary2:
        def __init__(self, a):
            self.a = a
        
        def __getitem__(self, b):
            return self.a + b
    
    def __getitem__(self, a):
        return Sparse2DArray.ary2(a)

a = Sparse2DArray()
print(a[1][2])

Comments

[goldberg]

ありがとうございます。このような書き方ができるとは、大変勉強になりました。