Cuando miras largo tiempo a un abismo, el abismo también mira dentro de ti
-- Friedrich Nietzsche
Si uno usa NumPy asiduamente lo normal es que se pase la vida consultando la extraordinaria documentación que hay en línea. Aunque está en un rinconcito discreto, de tanto ir el cántaro a la fuente, acaba uno por descubrir el enlace de [source] que aparece junto al nombre de la mayoría de funciones. No recuerdo cuándo fue la primera vez hice click en él, pero sí que recuerdo la sensación de desasosiego, de sentirme pequeñito en la inmensidad de una cantidad de código inabarcable y poco menos que incomprensible.
Probablemente uno de los módulos más agradecidos de visitar sea numpy.lib.arraysetops.py, que es donde viven funciones como numpy.unique , numpy.in1d y compañía. La gracia de este módulo viene del hecho de que implementa, usando sólo Python y unas cuántas funciones elementales de NumPy, las operaciones básicas con conjuntos.
La operación más elemental es obtener la lista de elementos únicos que hay en un array. La idea del método es bien sencilla: ordenar el array, quedarse con la primera ocurrencia de cada elemento. La implementación usa una máscara booleana, y viene a ser algo así como:
In [12]:
import numpy as np
In [13]:
a = np.random.randint(2, 6, size=8)
a
Out[13]:
In [14]:
a_sort = np.sort(a)
a_sort
Out[14]:
In [15]:
mask = a_sort[:-1] != a_sort[1:]
mask
Out[15]:
In [16]:
mask = np.concatenate(([True], mask))
mask
Out[16]:
In [17]:
unique = a_sort[mask]
unique
Out[17]:
Es todo razonablemente sencillo: comparamos cada elemento con el inmediatamente posterior, lo que nos permite descubrir todas las primeras ocurrencias de cada elemento único, excepto el primero, que lo añadimos explícitamente. Y aunque hay que ordenar el array, que es la operación que eventualmente domina por completo el tiempo de ejecución, para tamaños de array comedidos es incluso más rápido que implementaciones más elaboradas usando matrices asociativas (o sea, una hash table para que nos entendamos).
La cosa se complica un poco más cuando queremos deshacer la operación que hicimos al ordenar el array. Por poner un ejemplo concreto, np.unique se pueda llamar con la opción return_index=True, que devulve, además de los elementos únicos, el índice de la primera ocurrencia de cada uno en el array original. El truco está en ordenar el algoritmo indirectamente. Es decir, en vez de llamar a np.sort, llamamos a np.argsort para conseguir un array de índices que ordenen el array, índices que nos guardamos, y usamos acto seguido para ordenar el array.
Un detalle importante es que el algoritmo que usemos para ordenar el array tiene que ser estable. Es decir, que si hay elementos repetidos, su posición relativa en el array ordenado tiene que conservarse. El algoritmo por defecto de np.sort y np.argsort es quicksort, que no es estable. Afortunadamente tenemos la opción de usar mergesort que sí es estable. Teniendo esto en cuenta, el algoritmo completo vienes a ser tal que así:
In [18]:
a
Out[18]:
In [19]:
idx = np.argsort(a, kind='mergesort')
idx
Out[19]:
In [20]:
a_sort = a[idx]
a_sort
Out[20]:
In [21]:
mask = np.concatenate(([True], a_sort[:-1] != a_sort[1:]))
mask
Out[21]:
In [22]:
return_index = idx[mask]
return_index
Out[22]:
De nuevo es todo relativamente sencillo una vez que sabe uno cómo va la cosa, así que vamos a complicarlo todo un pelín más... A numpy.unique también se le puede pedir que return_inverse=True, o sea, que devuelva un array de índices que permitan reconstruir el array original a partir de los elementos únicos. Otra manera de decir lo mismo es que te devuelve un array, del mismo tamaño que el original, que contiene en cada posición, el índice del elemento único que se corresponde con esa posición. La cosa empieza de manera similar al caso anterior, construyendo idx (aunque el algoritmo no tiene por qué ser estable en este caso), a_sort y mask exactamente igual. Y a partir de aquí empieza la magia... Primero construimos el array de índices para el array ordenado, aprovechándonos de que al hacer operaciones aritméticas con valores booleanos, True se comporta como un 1, y False como un 0.
In [26]:
a
Out[26]:
In [23]:
return_inverse_sort = np.cumsum(mask) - 1
return_inverse_sort
Out[23]:
Ingenioso, ¿verdad? Ahora sólo queda desordenar este array para que recupere el ordenamiento original. Y eso, ¿cómo se hace? Pues la manera más sencilla, que es de hecho la que usan NumPy 1.9.x y anteriores, es la siguiente.
In [27]:
return_inverse = return_inverse_sort[np.argsort(idx)]
return_inverse
Out[27]:
No es evidente a primera vista, pero si se analiza con detenimiento no queda otra que asumir que sí, que efectivamente los índices que ordenan una lista de índices que ordenan un array, son los índices que devuelven al array ordenado el ordenamiento que tenía el array original. Cuesta casi tanto decirlo como aquello de "tres tristes tigres tragaban trigo en un trigal," pero tiene algo más de utilidad.
Hay un pequeño problema con hacerlo así: ordenar es de largo la operación más costosa de todo este procedimiento. Hemos tenido que hacerlo una vez, vale. ¿Pero es realmente necesario hacerlo la segunda? Pues no. Aunque hay que escribir más código, es mucho más eficiente hacerlo así:
In [29]:
return_inverse = np.empty_like(return_inverse_sort)
return_inverse[idx] = return_inverse_sort
return_inverse
Out[29]:
Uno de los cambios que vendráDe hecho, en NumPy 1.10.x