[Numpy-discussion] retrieving original array locations from 2d argsort

Warren Weckesser warren.weckesser at gmail.com
Mon Jul 15 21:23:30 EDT 2013


On 7/15/13, Moroney, Catherine M (398D)
<Catherine.M.Moroney at jpl.nasa.gov> wrote:
> I know that there's an easy way to solve this problem, but I'm not
> sufficiently knowledgeable
> about numpy indexing to figure it out.
>
> Here is the problem:
>
> Take a 2-d array a, of any size.
> Sort it in ascending order using, I presume, argsort.
> Step through the sorted array in order, and for each element in the sorted
> array,
> retrieve what the corresponding (line, sample) indices in the original array
> are.
>
> For instance:
>
> a = numpy.arange(0, 16).reshape(4,4)
> a[0,:] = -1*numpy.arange(0,4)
> a[2,:] = -1*numpy.arange(4,8)
>
> asort = numpy.sort(a, axis=None)
> for idx in xrange(0, asort.size):
> 	element = asort[idx]
>         !! Find the line and sample location in a that corresponds to the
> i-th element in assort
>


One way is to use argsort and  `numpy.unravel_index` to recover the
original 2D indices:

<code>
import numpy

a = numpy.arange(0, 16).reshape(4,4)
a[0,:] = -1*numpy.arange(0,4)
a[2,:] = -1*numpy.arange(4,8)

flat_sort_indices = numpy.argsort(a, axis=None)
original_indices = numpy.unravel_index(flat_sort_indices, a.shape)

print "  i   j  a[i,j]"
for i, j in zip(*original_indices):
        element = a[i,j]
        print "%3d %3d %6d" % (i, j, element)

</code>


Warren



> Thank-you for your help,
>
> Catherine
>
>
>
> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion at scipy.org
> http://mail.scipy.org/mailman/listinfo/numpy-discussion
>



More information about the NumPy-Discussion mailing list