连接范围数组以矢量化方式给出开始,停止数字 - NumPy

时间:2017-11-05 19:29:14

标签: python arrays numpy optimization vectorization

我有两个感兴趣的矩阵,第一个是"字袋"矩阵,有两列:文档ID和术语ID。例如:

bow[0:10]

Out[1]:
    array([[ 0, 10],
           [ 0, 12],
           [ 0, 19],
           [ 0, 20],
           [ 1,  9],
           [ 1, 24],
           [ 2, 33],
           [ 2, 34],
           [ 2, 35],
           [ 3, 2]])

另外,我有一个"索引"矩阵,其中矩阵中的每一行都包含单词矩阵包中给定文档ID的第一行和最后一行的索引。例如:第0行是doc id 0的第一个和最后一个索引。例如:

index[0:4]

Out[2]:
    array([[ 0,  4],
           [ 4,  6],
           [ 6,  9],
           [ 9, 10]])

我想要做的是随机抽取文档ID并获取这些文档ID的所有单词行。单词矩阵包大约150M行(~1.5Gb),因此使用numpy.in1d()太慢了。我们需要快速返回这些信息,以便为下游任务提供支持。

我提出的天真解决方案如下:

def get_rows(ids):
    indices = np.concatenate([np.arange(x1, x2) for x1,x2 in index[ids]])
    return bow[indices]

get_rows([4,10,3,5])

通用样本

提出问题的一般样本将是这样的 -

indices = np.array([[ 4, 7],
                    [10,16],
                    [11,18]]

预期的输出是 -

array([ 4,  5,  6, 10, 11, 12, 13, 14, 15, 11, 12, 13, 14, 15, 16, 17])

1 个答案:

答案 0 :(得分:4)

认为我最终使用cumsum技巧解决了它的矢量化解决方案 -

def create_ranges(a):
    l = a[:,1] - a[:,0]
    clens = l.cumsum()
    ids = np.ones(clens[-1],dtype=int)
    ids[0] = a[0,0]
    ids[clens[:-1]] = a[1:,0] - a[:-1,1]+1
    out = ids.cumsum()
    return out

样品运行 -

In [416]: a = np.array([[4,7],[10,16],[11,18]])

In [417]: create_ranges(a)
Out[417]: array([ 4,  5,  6, 10, 11, 12, 13, 14, 15, 11, 12, 13, 14, 15, 16, 17])

In [425]: a = np.array([[-2,4],[-5,2],[11,12]])

In [426]: create_ranges(a)
Out[426]: array([-2, -1,  0,  1,  2,  3, -5, -4, -3, -2, -1,  0,  1, 11])