Как извлечь значения из списка Python или массива NumPy, используя условные проверки с помощью применения NUMPY Векторизация?

Question

У меня есть следующий код, и я хочу извлечь определенные значения из других списков, которые зависят от данного условия. Но мои наборы данных огромны ~ 1 миллион значений в каждом списке. Поэтому этот метод вложенного цикла занимает слишком много времени. Есть ли векторизованный или более быстрый подход с использованием Numpy, который я могу использовать для ускорения своего кода, а также для использования меньшего количества памяти?

import random
import numpy as np

x=[random.randrange(0,10) for _ in range(0,100)]
y=[random.randrange(0,10) for _ in range(0,100)]
z=[random.randrange(0,10) for _ in range(0,100)]

x_unique=np.unique(x)

xx_list=[]
y_list=[]
z_list=[]

for i in range(len(x_unique)):
    xx_list.append([])
    y_list.append([])
    z_list.append([])

for ii, i in enumerate(x_unique):
        for j,k in enumerate(x):
            if i == k:
                xx_list[ii].append(x[j])
                y_list[ii].append(y[j])
                z_list[ii].append(z[j])

[РЕДАКТИРОВАТЬ: добавлен пример того, что ожидать]

В списках: y_list и z_list я хочу сохранить значения, которые соответствуют индексным номерам, которые хранятся в xx_list.

Например, рассмотрим следующие примеры списков:

x = [0.1,0.1,1,0.1,2,1,0.1]
y = [1.1,2.1,3,4,5,6,7]
z = [10,11,12,13.1,14,15,16]

Следовательно, x_unique имеет следующий вид:

x_unique = [0.1,1,2]

xx_list, y_list и z_list должны содержать следующее:

xx_list = [[0.1,0.1,0.1,0.1],[1,1],[2]]
y_list = [[1.1,2.1,4,7],[3,6],[5]]
z_list = [[10,11,13.1,16],[12,15],[14]]

Answer 1

Я нашел решение, которое занимает примерно 400 мс для 1M элементов, работающих со списками Python. И решение, которое занимает 100 мс при работе с массивами NumPy.

Python

Стратегия, которую я использую для построения одного словаря на каждый входной список (x, y, z). Каждый из них будет действовать как набор помеченных корзин. Для каждого входного списка bin i будет содержать элементы, для которых их соответствующий индекс в списке x равен i. Соответствующий означает, что они находятся на одной и той же позиции в соответствующем списке.

def compute_bins(x, y, z):
    # You can see this as an ordered-set:
    x_bin_indexes = {a:i for i, a in enumerate(sorted(set(x)))}

    # Each input list has its own set of labeled bins: 
    x_bins = defaultdict(list)
    y_bins = defaultdict(list)
    z_bins = defaultdict(list)

    for item_x, item_y, item_z in zip(x, y, z):
        index = x_bin_indexes[item_x]
        # Drop the item in the corresponding bin:
        x_bins[index].append(item_x)
        y_bins[index].append(item_y)
        z_bins[index].append(item_z)

    # Now we transform the result back to lists of list:
    x_bins = list(x_bins.values())
    y_bins = list(y_bins.values())
    z_bins = list(z_bins.values())
    return x_bins, y_bins, z_bins

Ключевым фактором здесь является то, что каждая операция, которую мы делаем в цикле, выполняется в постоянное время. Функцию можно вызвать так:

>>> xx_list, y_list, z_list = compute_bins(x, y, z)
>>> xx_list
[[0, 0, 0, 0], [1, 1], [2]]
>>> y_list
[[1, 2, 4, 7], [3, 6], [5]]
>>> z_list
[[10, 11, 13, 16], [12, 15], [14]]

Numpy

Используя numpy, я подумал о другой стратегии: отсортировать все массивы по элементам в x, а затем разделить их по количеству последовательных идентичных значений в x. Вот код (с учетом того, что x, y и z являются массивами numpy):

import numpy as np

def compute_bins(x, *others):
    x_bin_indexes, bin_sizes = np.unique(x, return_counts=True)
    sort_order = np.argsort(x)
    split_rule = np.cumsum(bin_sizes)[:-1]
    return tuple(np.split(o[sort_order], split_rule) for o in (x, *others))

Обратите внимание, что np.cumsum(bin_sizes)[:-1] существует только потому, что split принимает список индексов для вырезания, а не список размеров разрезов. Если мы хотим разделить [0, 0, 0, 1, 1, 2] на [[0, 0, 0], [1, 1], [2]], мы не передаем [3, 2, 1] в np.split, а вместо [3, 5].

Выступления

Что касается производительности, вот как это происходит на моей машине:

from random import randint

test_size = int(1e6)
x = [randint(0, 100) for _ in range(test_size)]
y = [i+1 for i in range(test_size)]
z = [i+test_size+1 for i in range(test_size)]

%timeit xx_list, y_list, z_list = compute_bins(x, y, z)

Выход для чистого python версия:

396 ms ± 5.98 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Выход для версии numpy (x, y и z являются np.array):

105 ms ± 1.07 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Для сравнения, первое предложенное вами решение дает:

19.7 s ± 282 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)