如何将[2,3,4]转换为[0,0,1,1,1,2,2,2,2]以利用tf.math.segment_sum？

Question

假设我有一个类似[2,3,4]的数组，我正在NumPy（或Tensorflow）中寻找一种将其转换为[0,0,1,1,1,2,2,2,2,2]的方法将tf.math.segment_sum（）应用于大小为2 + 3 + 4的张量。

我脑海中没有一个绝妙的主意，只有循环和列表理解。

Answer 1

这样的事情对您有用吗？

import numpy
arr = numpy.array([2, 3, 4])
numpy.repeat(numpy.arange(arr.size), arr)
# array([0, 0, 1, 1, 1, 2, 2, 2, 2])