Numpy: получить индекс элементов 1d массива в виде 2d массива

У меня есть такой массив: [1 2 2 0 0 1 3 5]

Можно ли получить индекс элементов в виде 2d массива? Например, ответ на приведенный выше ввод будет[[3 4], [0 5], [1 2], [6], [], [7]]

В настоящее время я должен зациклить различные значения и вызывать numpy.where(input == i)для каждого значения, которое имеет ужасную производительность с достаточно большим вводом.

Вот подход O (max (x) + len (x)) с использованием scipy.sparse:

import numpy as np
from scipy import sparse

x = np.array("1 2 2 0 0 1 3 5".split(),int)
x
# array([1, 2, 2, 0, 0, 1, 3, 5])


M,N = x.max()+1,x.size
sparse.csc_matrix((x,x,np.arange(N+1)),(M,N)).tolil().rows.tolist()
# [[3, 4], [0, 5], [1, 2], [6], [], [7]]

Это работает путем создания разреженной матрицы с записями в позициях (x [0], 0), (x [1], 1), ... Используя CSCформат (сжатый разреженный столбец), это довольно просто. Затем матрица преобразуется в LIL(связанный список) формат. Этот формат хранит индексы столбцов для каждой строки в виде списка в своемrows атрибуте, поэтому все, что нам нужно сделать, это взять и преобразовать его в список.

Обратите внимание, что для небольших массивов argsort решения, вероятно, быстрее, но при некоторых не слишком больших размерах это будет пересекаться.

РЕДАКТИРОВАТЬ:

argsortтолько numpyрешение:

np.split(x.argsort(kind="stable"),np.bincount(x)[:-1].cumsum())
# [array([3, 4]), array([0, 5]), array([1, 2]), array([6]), array([], dtype=int64), array([7])]

Если порядок индексов внутри групп не имеет значения, вы также можете попробовать argpartition(это не имеет значения в этом небольшом примере, но это не гарантируется в целом):

bb = np.bincount(x)[:-1].cumsum()
np.split(x.argpartition(bb),bb)
# [array([3, 4]), array([0, 5]), array([1, 2]), array([6]), array([], dtype=int64), array([7])]

РЕДАКТИРОВАТЬ:

@Divakar рекомендует против использования np.split. Вместо этого цикл, вероятно, быстрее:

A = x.argsort(kind="stable")
B = np.bincount(x+1).cumsum()
[A[B[i-1]:B[i]] for i in range(1,len(B))]

Или вы можете использовать совершенно новый (Python3.8 +) оператор моржа:

A = x.argsort(kind="stable")
B = np.bincount(x)
L = 0
[A[L:(L:=L+b)] for b in B.tolist()]

EDIT (отредактированный):

(Не чисто NumPy): В качестве альтернативы Numba (см. Пост @ Senderle) мы также можем использовать Pythran.

Компилировать с pythran -O3 <filename.py>

import numpy as np

#pythran export sort_to_bins(int[:],int)

def sort_to_bins(idx, mx):
    if mx==-1: 
        mx = idx.max() + 1
    cnts = np.zeros(mx + 2, int)
    for i in range(idx.size):
        cnts[idx[i] + 2] += 1
    for i in range(3, cnts.size):
        cnts[i] += cnts[i-1]
    res = np.empty_like(idx)
    for i in range(idx.size):
        res[cnts[idx[i]+1]] = i
        cnts[idx[i]+1] += 1
    return [res[cnts[i]:cnts[i+1]] for i in range(mx)]

Здесь numbaвыигрывает усиком по производительности:

repeat(lambda:enum_bins_numba_buffer(x),number=10)
# [0.6235917090671137, 0.6071486569708213, 0.6096088469494134]
repeat(lambda:sort_to_bins(x,-1),number=10)
# [0.6235359431011602, 0.6264424560358748, 0.6217901279451326]

Старые вещи:

import numpy as np

#pythran export bincollect(int[:])

def bincollect(a):
    o = [[] for _ in range(a.max()+1)]
    for i,j in enumerate(a):
        o[j].append(i)
    return o

Таймс против Нумба (старый)

timeit(lambda:bincollect(x),number=10)
# 3.5732191529823467
timeit(lambda:enumerate_bins(x),number=10)
# 6.7462647299980745

Один из возможных вариантов, в зависимости от размера ваших данных, это просто отказаться от numpyиспользования и использовать collections.defaultdict:

In [248]: from collections import defaultdict

In [249]: d = defaultdict(list)

In [250]: l = np.random.randint(0, 100, 100000)

In [251]: %%timeit
     ...: for k, v in enumerate(l):
     ...:     d[v].append(k)
     ...:
10 loops, best of 3: 22.8 ms per loop

Тогда вы получите словарь {value1: [index1, index2, ...], value2: [index3, index4, ...]}. Масштабирование времени довольно близко к линейному с размером массива, поэтому 10 000 000 занимает ~ 2,7 с на моей машине, что кажется достаточно разумным.

Хотя запрос на numpyрешение, я решил посмотреть, есть ли интересное numbaрешение. И действительно, есть! Вот подход, который представляет разделенный список как рваный массив, хранящийся в одном предварительно выделенном буфере. Это черпает вдохновение из argsortподхода, предложенного Полом Панцером . (Для более старой версии, которая не так хорошо, но была проще, см. Ниже.)

@numba.jit(numba.void(numba.int64[:], 
                      numba.int64[:], 
                      numba.int64[:]), 
           nopython=True)
def enum_bins_numba_buffer_inner(ints, bins, starts):
    for x in range(len(ints)):
        i = ints[x]
        bins[starts[i]] = x
        starts[i] += 1

@numba.jit(nopython=False)  # Not 100% sure this does anything...
def enum_bins_numba_buffer(ints):
    ends = np.bincount(ints).cumsum()
    starts = np.empty(ends.shape, dtype=np.int64)
    starts[1:] = ends[:-1]
    starts[0] = 0

    bins = np.empty(ints.shape, dtype=np.int64)
    enum_bins_numba_buffer_inner(ints, bins, starts)

    starts[1:] = ends[:-1]
    starts[0] = 0
    return [bins[s:e] for s, e in zip(starts, ends)]

Это обрабатывает список из десяти миллионов элементов за 75 мс, что почти в 50 раз быстрее по сравнению со списочной версией, написанной на чистом Python.

Для более медленной, но несколько более читаемой версии, вот что я имел раньше, основываясь на недавно добавленной экспериментальной поддержке «типизированных списков» с динамическим размером, которые позволяют нам гораздо быстрее заполнять каждую ячейку не по порядку.

Это numbaнемного борется с механизмом вывода типа, и я уверен, что есть лучший способ справиться с этой частью. Это также оказывается почти в 10 раз медленнее, чем выше.

@numba.jit(nopython=True)
def enum_bins_numba(ints):
    bins = numba.typed.List()
    for i in range(ints.max() + 1):
        inner = numba.typed.List()
        inner.append(0)  # An awkward way of forcing type inference.
        inner.pop()
        bins.append(inner)

    for x, i in enumerate(ints):
        bins[i].append(x)

    return bins

Я проверил это в отношении следующего:

def enum_bins_dict(ints):
    enum_bins = defaultdict(list)
    for k, v in enumerate(ints):
        enum_bins[v].append(k)
    return enum_bins

def enum_bins_list(ints):
    enum_bins = [[] for i in range(ints.max() + 1)]
    for x, i in enumerate(ints):
        enum_bins[i].append(x)
    return enum_bins

def enum_bins_sparse(ints):
    M, N = ints.max() + 1, ints.size
    return sparse.csc_matrix((ints, ints, np.arange(N + 1)),
                             (M, N)).tolil().rows.tolist()

Я также проверил их на предварительно скомпилированной версии Cython, аналогичной enum_bins_numba_buffer (подробно описанной ниже).

В списке из десяти миллионов случайных чисел ( ints = np.random.randint(0, 100, 10000000)) я получаю следующие результаты:

enum_bins_dict(ints)
3.71 s ± 80.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_list(ints)
3.28 s ± 52.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_sparse(ints)
1.02 s ± 34.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_numba(ints)
693 ms ± 5.81 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_cython(ints)
82.3 ms ± 1.77 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

enum_bins_numba_buffer(ints)
77.4 ms ± 2.06 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Впечатляет, что этот способ работы с numbaопережает cythonверсию той же функции, даже с отключенной проверкой границ. У меня пока нет достаточных знаний, pythranчтобы протестировать этот подход, используя его, но мне было бы интересно увидеть сравнение. Вероятно, исходя из этого ускорения,pythran версия также может быть немного быстрее при таком подходе.

Вот cythonверсия для справки, с некоторыми инструкциями по сборке. После cythonустановки вам понадобится простой setup.pyфайл, подобный следующему:

from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize
import numpy

ext_modules = [
    Extension(
        'enum_bins_cython',
        ['enum_bins_cython.pyx'],
    )
]

setup(
    ext_modules=cythonize(ext_modules),
    include_dirs=[numpy.get_include()]
)

И модуль Cython enum_bins_cython.pyx:

# cython: language_level=3

import cython
import numpy
cimport numpy

@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef void enum_bins_inner(long[:] ints, long[:] bins, long[:] starts) nogil:
    cdef long i, x
    for x in range(len(ints)):
        i = ints[x]
        bins[starts[i]] = x
        starts[i] = starts[i] + 1

def enum_bins_cython(ints):
    assert (ints >= 0).all()
    # There might be a way to avoid storing two offset arrays and
    # save memory, but `enum_bins_inner` modifies the input, and
    # having separate lists of starts and ends is convenient for
    # the final partition stage.
    ends = numpy.bincount(ints).cumsum()
    starts = numpy.empty(ends.shape, dtype=numpy.int64)
    starts[1:] = ends[:-1]
    starts[0] = 0

    bins = numpy.empty(ints.shape, dtype=numpy.int64)
    enum_bins_inner(ints, bins, starts)

    starts[1:] = ends[:-1]
    starts[0] = 0
    return [bins[s:e] for s, e in zip(starts, ends)]

С этими двумя файлами в вашем рабочем каталоге выполните эту команду:

python setup.py build_ext --inplace

Затем вы можете импортировать функцию, используя from enum_bins_cython import enum_bins_cython.

Вот действительно очень странный способ сделать это, это ужасно, но я нашел это слишком смешным, чтобы не делиться - и все numpy!

out = np.array([''] * (x.max() + 1), dtype = object)
np.add.at(out, x, ["{} ".format(i) for i in range(x.size)])
[[int(i) for i in o.split()] for o in out]

Out[]:
[[3, 4], [0, 5], [1, 2], [6], [], [7]]

РЕДАКТИРОВАТЬ: это лучший метод, который я мог найти на этом пути. Это все еще в 10 раз медленнее, чем решение @PaulPanzer argsort:

out = np.empty((x.max() + 1), dtype = object)
out[:] = [[]] * (x.max() + 1)
coords = np.empty(x.size, dtype = object)
coords[:] = [[i] for i in range(x.size)]
np.add.at(out, x, coords)
list(out)

Вы можете сделать это, составив словарь чисел, ключи - это числа, а значения - это индексы, которые видели числа, это один из самых быстрых способов сделать это, вы можете увидеть код ниже:

>>> import numpy as np
>>> a = np.array([1 ,2 ,2 ,0 ,0 ,1 ,3, 5])
>>> b = {}
# Creating an empty list for the numbers that exist in array a
>>> for i in range(np.min(a),np.max(a)+1):
    b[str(i)] = []

# Adding indices to the corresponding key
>>> for i in range(len(a)):
    b[str(a[i])].append(i)

# Resulting Dictionary
>>> b
{'0': [3, 4], '1': [0, 5], '2': [1, 2], '3': [6], '4': [], '5': [7]}

# Printing the result in the way you wanted.
>>> for i in sorted (b.keys()) :
     print(b[i], end = " ")

[3, 4] [0, 5] [1, 2] [6] [] [7] 

псевдокод:

1. получите «количество 1d массивов в 2d массиве», вычитая минимальное значение вашего массива numpy из максимального значения, а затем плюс один. В вашем случае это будет 5-0 + 1 = 6

2. инициализировать 2d массив с количеством 1d массивов в нем. В вашем случае инициализируйте 2d массив с 6 1d массивом в нем. Каждый 1d массив соответствует уникальному элементу в вашем массиве numpy, например, первый 1d массив будет соответствовать '0', второй 1d массив будет соответствовать '1', ...

3. переберите ваш массивный массив, поместите индекс элемента в соответствующий соответствующий 1d массив. В вашем случае индекс первого элемента в вашем массиве numpy будет помещен во второй массив 1d, индекс второго элемента в вашем массиве numpy будет помещен в третий массив 1d, ....

Этот псевдокод будет работать линейно, так как он зависит от длины вашего массива.

Это дает вам именно то, что вы хотите, и заняло бы около 2,5 секунд для 10 000 000 на моей машине:

import numpy as np
import timeit

# x = np.array("1 2 2 0 0 1 3 5".split(),int)
x = np.random.randint(0, 100, 100000)

def create_index_list(x):
    d = {}
    max_value = -1
    for i,v in enumerate(x):
        if v > max_value:
            max_value = v
        try:
            d[v].append(i)
        except:
            d[v] = [i]
    result_list = []
    for i in range(max_value+1):
        if i in d:
            result_list.append(d[i])
        else:
            result_list.append([])
    return result_list

# print(create_index_list(x))
print(timeit.timeit(stmt='create_index_list(x)', number=1, globals=globals()))

Итак, учитывая список элементов, вы хотите составить (элемент, индекс) пары. В линейное время это можно сделать так:

hashtable = dict()
for idx, val in enumerate(mylist):
    if val not in hashtable.keys():
         hashtable[val] = list()
    hashtable[val].append(idx)
newlist = sorted(hashtable.values())

Это должно занять O (N) время. Я не могу придумать более быстрого решения на данный момент, но обновлю здесь, если я сделаю.