Numpy 3Darrayのインデックスを2Darrayで指定したい(再度)

Question

https://teratail.com/questions/172545
にて、np.chooseを使えばよいのではとご回答いただき、試したところchooseは要素数が32までに制限されるようで、再質問させていただきます。前回回答くださった方には、質問の情報が不十分で申し訳ありませんでした。

前提・実現したいこと

(n×m×l)の3次元の配列Aでそれぞれx軸、y軸、z軸と考えます。
z軸のインデックスを与える(n×m)の配列Bを使って、指定するz軸の値を切り出した(n×m)の配列aを得たいです。実際のデータのn,m,lはそれぞれ8000,30,200程度。
つまり以下のコードでやっているようなことです。

import numpy as np
import itertools

A = np.arange(2*3*4).reshape(2,3,4) #元の三次元配列
B = np.array([[0,1,1],[1,2,3]])     #Z軸のインデックスを示す配列
a = np.zeros((2,3))　　　#Aから指定するZ軸の値を取りだして格納する配列

for i,j in itertools.product([0,1],[0,1,2]):
    a[i,j] = A[i,j,B[i,j]]

試したこと

numpy.chooseを使う→chooseの要素数の制限でだめ。
Bが示すz軸インデックスの要素のみ1で他は0のAと同じshapeの配列bを作ってA[b==1]で取り出す→一応意図通り動くのだがスマートではないような…
numpy.selectを使うと実現できそうな気がするのですが、conditionをどう書けばよいのかわかりませんでした。

Pythonらしい解決法をご教示いただければ、ありがたいです。よろしくお願いいたします。

Answer 1

スマートではないけど。

np.take_along_axis(A, B.reshape([*B.shape, 1]), axis=2).reshape(B.shape)

Comments

[退会済みユーザー]

ご回答ありがとうございます。numpy1.14で止まっていたので、アップデートしてtake_along_axisが使えるようにしました。
実行しようとするとB.reshape([*B.shape, 1])の「*」にinvalid syntax エラーが出るのですが、何をしているどういう文法ですか？不勉強で、初めて見ました。Python2.7でも使えるのでしょうか。

[hayataka2049]

こういう文法です。
>>> [*[1,2,3], 4]
[1, 2, 3, 4]
python3.5以降のはずです。
代替したければB.shape+(1,)でいいはずです（そっちの方が素直かも）

[退会済みユーザー]

*のついている配列の中身をバラす、みたいな感じでしょうか。3系に切り替えたら使ってみます。
やっているのはBをm×n×1にreshapeするという操作ですね。B.shape+(1,)に置き換えて、できました。ありがとうございました！
PS いつもブログで勉強させていただいてます。