Search code examples
pythonnumpymatrixenumeration

Enumerating python matrix by cell value


Given a matrix:

m = np.matrix('1 2; 3 4')

I would like to enumerate its cells (not rows\columns) in descending order of the value in the matrix. So in my example i would like to enumerate: 4,3,2,1 and have the indexes for each.

I could flatten the matrix to 1d array, order the array and use its index to recover to the original matrix index. But it feels wrong.

Is there some built-in ability in python for that?


Solution

  • You could use np.ndenumerate to get an iterator of indices and values.

    You could convert it to a list:

    >>> import numpy as np
    >>> m = np.matrix('1 2; 3 4')
    >>> list(np.ndenumerate(m))
    [((0, 0), 1), ((0, 1), 2), ((1, 0), 3), ((1, 1), 4)]
    

    sort by values:

    >>> sorted([(v, i) for (i,v) in np.ndenumerate(m)], reverse=True)
    [(4, (1, 1)), (3, (1, 0)), (2, (0, 1)), (1, (0, 0))]
    

    and extract the indices:

    >>> [i for (_, i) in sorted([(v, i) for (i,v) in np.ndenumerate(m)], reverse=True)]
    [(1, 1), (1, 0), (0, 1), (0, 0)]
    

    I don't know if it's possible to do it in less steps. Here's an example with another matrix:

    >>> m = np.matrix('1 4; 3 2')
    >>> [i for (_, i) in sorted([(v, i) for (i,v) in np.ndenumerate(m)], reverse=True)]
    [(0, 1), (1, 0), (1, 1), (0, 0)]
    

    It might seem messy to mix numpy and vanilla list comprehensions. The problem is that ndenumerate returns an iterator. You could use:

    np.fromiter(np.ndenumerate(m), np.dtype('2i, 1i'))
    

    but it's a 1-D array of tuples, which doesn't makes indexing or sorting easier.