字典：按多个键+降序/升序排序

Question

I have a dict with keys as str and values as np.array . The types of the np.array can also be of type np.str_ .

data = {"col1": np.array((1, 2, 3, 4, 5, 4, 3, 2, 1)),
        "col2": np.array(list("abcdeabcd")),
        "col3": np.array((10, 11, 9, 8, 7, 2, 12, 100, 1))}

How can I sort by multiple keys and ascending/descending order similar to what I could do with pandas 's sort_values method:

Pandas solution (not needed)

df = pd.DataFrame(data)
df.sort_values(by=["col1", "col2"], ascending=[True, True])

Numpy or base python solution needed:

I don't want to use pandas but ideally something in numpy . I know that I can use np.lexsort to sort by multiple columns. But this does not give me (i) the option to do ascending/descending.

This is the output I want when sorting by col1 and then col2 in ascending / ascending order:

{'col1': np.array([1, 1, 2, 2, 3, 3, 4, 4, 5]),
 'col2': np.array(['a', 'd', 'b', 'c', 'b', 'c', 'a', 'd', 'e']),
 'col3': np.array([10, 1, 11, 100, 12, 9, 2, 8, 7])}

This is the output I want when sorting by col1 and then col2 in ascending / descending order:

{'col1': np.array([1, 1, 2, 2, 3, 3, 4, 4, 5]),
 'col2': np.array(['d', 'a', 'c', 'b', 'c', 'b', 'd', 'a', 'e']),
 'col3': np.array([1, 10, 100, 11, 9, 12, 8, 2, 7])}

Answer 1

You can use np.lexsort to sort columns backwards by negating the sort key. Remember that you don't need to pass the actual array to lexsort . For signed numerical arrays, you can sort in reverse by negating the values. For unsigned integers, you can subtract the values from the maximum. Strings can either be treated as numbers of you can make a lookup table of signed integers based on np.unique .

Here is a small example of arrays:

np.random.seed(0xBEEF)
a = np.random.choice([1, 2, 3], 10)
b = np.random.choice([1.0, 2.0, 3.0], 10)
c = np.random.choice(np.array([1, 2, 3], dtype=np.uint8), 10)
d = np.random.choice(list('abc'), 10)

The sort keys in ascending order can be the arrays themselves in all cases. In descending order, we can obviously use -a and -b . As it happens, -c also works:

>>> c
array([3, 1, 1, 3, 2, 2, 1, 1, 2, 3], dtype=uint8)
>>> -c
array([253, 255, 255, 253, 254, 254, 255, 255, 254, 253]

This may be dependent on the platform representation, but on most popular systems, negative numbers are represented in two's complement form and this should work just fine. If you wanted to be really safe, you could add a check like

if np.issubdtype(c.dtype, np.unsignedinteger):
    key = np.iinfo(c.dtype).max + 1 - c

And of course we have d :

>>> d
array(['c', 'a', 'b', 'b', 'c', 'b', 'a', 'b', 'c', 'b'], dtype='<U1')
>>> -d
...
UFuncTypeError: ufunc 'negative' did not contain a loop with signature matching types dtype('<U1') -> dtype('<U1')

One way to construct a sort key here is:

lookup, key = np.unique(d, return_inverse=True)

The elements of key are indices into lookup , which is in sorted order, meaning that if you sorted key , the result of lookup[key] would be correctly sorted as well. This means that key.argsort() and d.argsort() are the same, with the added advantage that you can negate key .

In fact, you can take a shortcut an write your key generator using this technique alone:

def make_key(arr, asc=True):
    _, key = np.unique(arr, return_inverse=True)
    if not asc:
        key = np.negative(key, out=key) # Don't bother making a second array
    return key

So your full example could look something like this:

def custom_lexsort(arrs, asc=True):
    """
    Lexsort a collection of arrays in ascending or descending order.

    Parameters
    ----------
    arrs : sequence[array-like]
        Sequence of arrays to sort.
    asc : array-like[bool]
        Sequence of True for ascending elements of `keys`,
        False for descending. Must broadcast to `(len(arrs),)`.
    """
    def make_key(a, asc):
        if np.issubdtype(a.dtype, np.number):
            key = a
        else:
            _, key = np.unique(a, return_inverse=True)
        if asc:
            return key
        elif np.issubdtype(key.dtype, np.unsignedinteger):
            return np.iinfo(key.dtype).max + 1 - key
        else:
            return -key

    n = len(arrs)
    asc = np.broadcast_to(asc, n)
    keys = [make_key(*x) for x in zip(arrs, asc)]
    return np.lexsort(keys[::-1])

data = {"col1": np.array((1, 2, 3, 4, 5, 4, 3, 2, 1)),
        "col2": np.array(list("abcdeabcd")),
        "col3": np.array((10, 11, 9, 8, 7, 2, 12, 100, 1))}

idx = custom_lexsort(list(data.values()), [True, False, True])
result = {k: v[idx] for k, v in data.items()}

I've taken the liberty of reversing the order of the arrays, since lexsort sorts from last to first. And sure enough:

>>> result
{'col1': array([1, 1, 2, 2, 3, 3, 4, 4, 5]),
 'col2': array(['d', 'a', 'c', 'b', 'c', 'b', 'd', 'a', 'e'], dtype='<U1'),
 'col3': array([  1,  10, 100,  11,   9,  12,   8,   2,   7])}

I've included the third column for sorting because it does no harm. Here is an example of the arrays sorted in place, with only the first two used for sorting:

idx = custom_lexsort([data['col1'], data['col2']], [True, False])
for v in data.values():
    v[:] = v[idx]