numpy的bincount函数的逆操作

5 投票
3 回答
1296 浏览
提问于 2025-04-18 00:02

给定一个整数数组 c,我该如何把它转换成另一个整数数组 inds,使得 np.all(np.bincount(inds) == c) 这个条件成立呢?

举个例子:

>>> c = np.array([1,3,2,2])
>>> inverse_bincount(c)  # <-- what I need

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

背景:我正在尝试跟踪多个数据集的位置,同时对它们进行一次性计算。我把所有数据拼接在一起进行批处理,但我需要一个索引数组来提取结果。

目前的解决方法:

def inverse_bincount(c):
  return np.array(list(chain.from_iterable([i]*n for i,n in enumerate(c))))

3 个回答

1

下面这个方法在我的电脑上运行速度大约是目前被接受的答案的两倍;不过我得说,我对np.repeat的表现感到很惊讶。我本以为它会因为频繁创建临时对象而变得很慢,但实际上它的表现相当不错。

import numpy as np
c = np.array([1,3,2,2])
p = np.cumsum(c)
i = np.zeros(p[-1],np.int)
np.add.at(i, p[:-1], 1)
print np.cumsum(i)
1

这里说的是不需要用到numpy这个库:

c = [1,3,2,2]
reduce(lambda x,y: x + [y] * c[y], range(len(c)), [])
13

使用 numpy.repeat :

np.repeat(np.arange(c.size), c)

撰写回答