如何从numpy concatenate获取索引

时间:2017-09-21 01:52:47

标签: python arrays numpy indexing concatenation

我希望concatenate numpy ndarray s,但也有索引。

也就是说,

>>> np.concatenate([[1,2,3,4], [11,22,33], [4], [435,456]])

我希望结果numpy.concatenate已经给出:

array([  1,   2,   3,   4,  11,  22,  33,   4, 435, 456])

但我希望输入列表的索引指示输出列表的哪个元素是输出数组的每个元素来自:

array([0, 0, 0, 0, 1, 1, 1, 2, 3, 3]

我知道基于np.cumsumnp.repeat的简单配方,我可以用它来自己构建索引数组,但是有一种简单的本地方式吗?

2 个答案:

答案 0 :(得分:2)

这是一个基于直接列表的解决方案:

In [64]: alist = [[1,2,3,4], [11,22,33], [4], [435,456]]
In [67]: idlist = [[i+1]*len(x) for i,x in enumerate(alist)]
In [68]: np.concatenate(alist)
Out[68]: array([  1,   2,   3,   4,  11,  22,  33,   4, 435, 456])
In [69]: np.concatenate(idlist)
Out[69]: array([1, 1, 1, 1, 2, 2, 2, 3, 4, 4])

如果子列表非常长,或者已经是数组,那么使用更多numpy操作的解决方案可能会更快。但如果这是典型的子列表组合,列表理解可能会更快。请记住,concatenate的输入是一个列表(即使输入是一个数组,它也被视为一个列表。)

In [73]: timeit idlist = [np.arange(len(x)) for i,x in enumerate(alist)]
6.06 µs ± 224 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [74]: timeit idlist = [[i+1]*len(x) for i,x in enumerate(alist)]
2.98 µs ± 8.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

答案 1 :(得分:0)

这个怎么样?

In [16]: alist = [[1,2,3,4], [11,22,33], [4], [435,456]]

In [17]: my_arrays = [np.array(x) for x in alist]

In [18]: index_arrays = [np.ones((x.shape), dtype=int)*i for i, x in enumerate(my_arrays)]

In [19]: np.concatenate(my_arrays)
Out[19]: array([  1,   2,   3,   4,  11,  22,  33,   4, 435, 456])

In [20]: np.concatenate(index_arrays)
Out[20]: array([0, 0, 0, 0, 1, 1, 1, 2, 3, 3])