Как преобразовать [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], чтобы применить 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])

Answer 2

Вам не нужно использовать NumPy. Вы можете использовать только списки:

>>> foo = [2,3,4]
>>> sum([[i]*foo[i] for i in range(len(foo))], [])
[0, 0, 1, 1, 1, 2, 2, 2, 2]

Это работает так:

Вы можете создавать расширенные массивы, умножая простой массив на константу, поэтому [0] * 2 == [0,0]. Таким образом, для каждого индекса в массиве мы расширяемся на [i]*foo[i]. Другими словами:

>>> [[i]*foo[i] for i in range(len(foo))]
[[0, 0], [1, 1, 1], [2, 2, 2, 2]]

Затем мы используем sum, чтобы свести списки в один список:

>>> sum([[i]*foo[i] for i in range(len(foo))], [])
[0, 0, 1, 1, 1, 2, 2, 2, 2]

Поскольку мы «суммируем» списки, а не целые числа, мы передаем[] до sum, чтобы сделать пустой список начальным значением суммы.

(Обратите внимание, что это, скорее всего, будет медленнее, чем numpy, хотя лично я не сравнивал его с чем-то вроде ответа @ Patol75. )

Answer 3

Мне очень нравится ответ от @ Patol75, так как он аккуратный. Тем не менее, пока не существует чистого решения с тензорным потоком, поэтому я предлагаю решение, которое может быть довольно сложным. Просто для справки и для удовольствия!

Кстати, я не видел tf.repeat этого API в tf master. Пожалуйста, отметьте PR , который добавляет tf.repeat поддержку, эквивалентную numpy.repeat.

import tensorflow as tf

repeats = tf.constant([2,3,4])
values = tf.range(tf.size(repeats))  # [0,1,2]

max_repeats = tf.reduce_max(repeats)  # max repeat is 4

tiled = tf.tile(tf.reshape(values, [-1,1]), [1,max_repeats])  # [[0,0,0,0],[1,1,1,1],[2,2,2,2]]

mask = tf.sequence_mask(repeats, max_repeats)  # [[1,1,0,0],[1,1,1,0],[1,1,1,1]]

res = tf.boolean_mask(tiled, mask)  # [0,0,1,1,1,2,2,2,2]

Answer 4

В ответе Patol75 используется Numpy, но Gort the Robot's answer на самом деле быстрее (по крайней мере, в вашем списке примеров).

Я буду держать этот ответ какдругое решение, но оно медленнее, чем оба.

Учитывая, что a = [2,3,4] это можно сделать, используя цикл, подобный следующему:

b = []

for i in range(len(a)):
   for j in range(a[i]):
      b.append(range(len(a))[i])

, который, как однострочник для понимания списка,эта дьявольская вещь:

b = [range(len(a))[i] for i in range(len(a)) for j in range(a[i])]

Оба заканчиваются b = [0,0,1,1,1,2,2,2,2].