import tensorflow as tf
arr=tf.constant([1, 5, 7, 4, 3, 1, 9])
searched_value=tf.constant([1, 5, 9])
index_arr=tf.where(tf.equal(arr,tf.expand_dims(bbb,axis=-1)))
'''
Out:
array([[0, 0],
[0, 5],
[1, 1],
[2, 6]], dtype=int64)
'''
final_index=tf.slice(index_arr,[0,1],[-1,1])#获取索引为1的列
'''
Out:
array([[0],
[5],
[1],
[6]], dtype=int64)
'''
final_index=tf.reshape(final_index,[-1])#展平tensor,[0,5,1,6]
final_value=tf.gather(arr,final_index)#
'''
Out:
[1,1,5,9]
'''
其中,采用了Broadcasting机制,详情查看https://docs.scipy.org/doc/numpy-1.13.0/user/basics.broadcasting.html
在多维数组中查找多个元素时,如上例子:
arr维度7,searched_value维度3,通过将searched_value在最后插入1维,使searched_value维度变为(3,1)
通过broadcasting机制,实行维度对齐:
arr: 7 1×7 3×7
->补充维度 ->复制
searched_value: 3×1 3×1 3×7
最终:
arr=
[[1, 5, 7, 4, 3, 1, 9],
[1, 5, 7, 4, 3, 1, 9],
[1, 5, 7, 4, 3, 1, 9]]
searched_value=
[[1,1,1,1,1,1,1],
[5,5,5,5,5,5,5],
[9,9,9,9,9,9,9]]
接着tf.equal(arr,searched_value)->[[True,False,False,False,False,True,False],
[False,True,False,False,False,False,False],
[False,False,False,False,False,False,True]]
tf.where将True对应的索引输出,行索引[0,0,1,2],列索引[0,5,1,6]