再帰的 zip

はじめに

numpy.ndarrayには ファンシーインデックス という仕組みがある．大変便利であるが，インデックスの入れ子構造が私の直感と逆で使いにくい．例として，先の記事から次のコードを引用しよう．

import numpy as np

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
a[[0, 1], [2, 2]] #=> [3, 6]

これ，[0, 2]要素と[1, 2]要素を抜き出すためのインデックスが ([0, 1], [2, 2]) なのであるが，私は a[[0, 2], [1, 2]] と書きたい (まあ，多分ブロードキャストやスライスと併用したいからこの入れ子順なんだろうけど，私は直感的に書きたいのだ)．このような入れ子の入れ替えには zip 関数 が利用できる：

a[tuple(zip([0, 2], [1, 2]))] #=> array([3, 6])
# zip オブジェクトは直接インデックスにできないので，tuple をかけている

しかし，numpy.ndarrayはより高階のテンソルである場合もあり，単純なzipでは済まないこともあるので，再帰的に入れ替えていくような関数がほしい．

実装

要件を言語化するのが難しいので，試行錯誤の結果を先に示そう．

import collections

def deepzip(*x):
  if isinstance(x[0][0], collections.abc.Iterable):
    return zip(*map(lambda y: deepzip(*y), x))
  else:
    return zip(*x)
def invert_index(idx):
  return tuple(deepzip(*idx))

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
a[invert_index(([[0, 2], [1, 2]], [[1, 0], [1, 1]]))] #=> array([[3, 6], [4, 5]])

Iterableかどうかで判定しているので，インデックスがリストじゃなくてタプルやndarrayでもOK．x[0][0]はキモいが，深さが揃ってないとかの不正処理は省略した．入力，出力ともに (整合していれば) 好きな階数でいける……はず．

上手く行かないパターンとか，こうした方がいいとかあったら教えて下さい．