如何在另一个张量中找到张量 A 元素的索引？

Question

>假设我们有

a = tf.constant(['a', 'c'])
b = tf.constant(['a', 'd', 'c'])
>>> Intended output: [0, 2] since b[0] = a[0] and b[2] = a[1]

如何找到 a 在 b 中的位置？

我认为所需的操作是tf.gather()的反向/反向。

我知道一个非常占用空间的解决方案：

tf.where(tf.equal(tf.expand_dims(a, 1), tf.expand_dims(b, 0)))[:, 1]  >> array([0, 2])

但我认为tf.equal()会产生形状的张量(2,3(。它占用太多空间。(空格 O(m * n( 在我的情况下太多了。

有没有更好的解决方案来节省空间？

Answer 1

要用时间换取空间：

result = tf.where(tf.equal(b, a[0]))
t = tf.constant(1)
cond = lambda t, result: tf.less(t, tf.shape(a)[0])
body = lambda t, result: [t + 1, tf.concat([result, tf.where(tf.equal(b, a[t]))], axis=0)]
result = tf.while_loop(cond, body, [t, result], shape_invariants=[t.get_shape(), tf.TensorShape([None, None])])[1]
sess.run(tf.squeeze(result))

如果张量 A 的大小为 m，张量 b 的大小为 n，则此解循环 m 次，占用 O(m*n( 时间和 O(m+n( 空间。