从概率质量函数的二维数组的每列中抽样随机索引

Question

我有一个名为pmf的二维数组，其中每列都是概率质量函数。我想根据相应的概率质量函数对每列采样一行索引。我目前有

[np.random.choice(range(len(pmf.T[i])), 1, p=pmf.T[i])[0] for i in xrange(len(pmf.T))]

是否有更清洁的numpy方式？

举个例子：

pmf = [[0.1215122141454291, 0.02526931631702081], 
       [0.44580918821894255, 0.23465429357912862],
       [0.13924078748403307, 0.7311930810111874], 
       [0.2934378101515954, 0.008883309092663008]]

然后输出可能是：

[3, 2]