CODEWITHSUNDEEP
pythontensorflow
SameeraR
SameeraR

Reputation: 415

Sort a tensor based on two columns in tensorflow

Is it possible to sort a tensor based on values in two columns in Tensorflow?

For example, let's say I have the following tensor.

[[1,2,3]
[2,3,5]
[1,4,6]
[2,2,1]
[0,4,2]]

I would lik it to be sorted first based on the first column and then the second column. After sorting it would look like following.

[[0,4,2]
[1,2,3]
[1,4,6]
[2,2,1]
[2,3,5]]

Is there any way to achieve this using tensorflow? I am able to sort according to a single column. But sorting based on two columns is a problem for me. Please can anyone help?

Upvotes: 2

Views: 1486

Answers (1)

Nipun Wijerathne
Nipun Wijerathne

Reputation: 1829

A very naive approach, 

import tensorflow as tf

a = tf.constant([[1, 2, 3],
                 [2, 3, 5],
                 [1, 4, 6],
                 [2, 2, 1],
                 [0, 4, 2]])

# b = a[:0]*10 + a[:1]*1 -- > (e.g 1*10+2*1 =12)
b = tf.add(tf.slice(a, [0, 0], [-1, 1]) * 10, tf.slice(a, [0, 1], [-1, 1]))

reordered = tf.gather(a, tf.nn.top_k(b[:, 0], k=5, sorted=False).indices)
reordered = tf.reverse(reordered, axis=[0])

with tf.Session() as sess:
    result = sess.run(reordered)
    print(result)

Hope this helps.

Upvotes: 1

Related Questions