python - 使用列表中的三元组索引 numpy 数组

Question

我有一个表示三元语言模型的 3d numpy 数组 A。所以 A[i, i-1, i-2] 是概率 $P(w_i|w_{i-1},w_{i-2})$，其中 $w$ 是连续的单词。我想提取一系列单词的所有概率。现在我正在使用以下内容：

probas = [] 
for i in range(2, len(words)):
    probas.append(A[words[i-2], words[i-2], words[i]])

词是某些词汇中词的索引。我的问题是，这可以通过 numpy 索引有效地矢量化吗？单词列表可能真的很长，并且以这种方式循环可能会非常昂贵，因为它将针对许多序列进行。