Numpy: lấy chỉ mục của các phần tử của mảng 1d dưới dạng mảng 2d

Tôi có một mảng numpy như thế này: [1 2 2 0 0 1 3 5]

Có thể lấy chỉ mục của các phần tử dưới dạng mảng 2d không? Ví dụ, câu trả lời cho đầu vào trên sẽ là[[3 4], [0 5], [1 2], [6], [], [7]]

Hiện tại tôi phải lặp các giá trị khác nhau và gọi numpy.where(input == i)cho từng giá trị, có hiệu suất khủng với đầu vào đủ lớn.

Đây là cách tiếp cận O (max (x) + len (x)) bằng cách sử dụng 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]]

Điều này hoạt động bằng cách tạo một ma trận thưa thớt với các mục tại các vị trí (x [0], 0), (x [1], 1), ... Sử dụng CSCđịnh dạng (cột thưa nén) điều này khá đơn giản. Ma trận sau đó được chuyển đổi sang LILđịnh dạng (danh sách liên kết). Định dạng này lưu trữ các chỉ mục cột cho mỗi hàng dưới dạng một danh sách trong rowsthuộc tính của nó , vì vậy tất cả những gì chúng ta cần làm là lấy đó và chuyển đổi nó thành danh sách.

Lưu ý rằng đối với các argsortgiải pháp dựa trên mảng nhỏ có thể nhanh hơn nhưng ở một số kích thước không lớn, điều này sẽ vượt qua.

BIÊN TẬP:

argsortnumpygiải pháp dựa trên cơ sở :

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])]

Nếu thứ tự các chỉ số trong các nhóm không thành vấn đề, bạn cũng có thể thử argpartition(điều này xảy ra không có sự khác biệt trong ví dụ nhỏ này nhưng nói chung điều này không được đảm bảo):

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])]

BIÊN TẬP:

@Divakar khuyến cáo không nên sử dụng np.split. Thay vào đó, một vòng lặp có thể nhanh hơn:

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

Hoặc bạn có thể sử dụng toán tử walrus hoàn toàn mới (Python3.8 +):

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

EDIT (EDITED):

(Không thuần túy numpy): Thay thế cho numba (xem bài đăng của @ senderle) chúng ta cũng có thể sử dụng pythran.

Biên dịch với 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)]

Ở đây numbachiến thắng bởi một hiệu suất khôn ngoan:

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]

Những thứ cũ hơn:

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

Thời gian so với numba (cũ)

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

Một tùy chọn tiềm năng tùy thuộc vào kích thước dữ liệu của bạn là bỏ ra numpyvà sử dụng 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

Sau đó, bạn kết thúc với một từ điển của {value1: [index1, index2, ...], value2: [index3, index4, ...]}. Tỷ lệ thời gian khá gần với tuyến tính với kích thước của mảng, vì vậy 10.000.000 mất ~ 2.7 giây trên máy của tôi, có vẻ đủ hợp lý.

Mặc dù yêu cầu là một numpygiải pháp, tôi quyết định xem liệu có một numbagiải pháp dựa trên thú vị nào không . Và thực sự là có! Đây là một cách tiếp cận đại diện cho danh sách được phân vùng dưới dạng một mảng rách rưới được lưu trữ trong một bộ đệm preallocated. Điều này lấy một số cảm hứng từ argsortcách tiếp cận được đề xuất bởi Paul Panzer . (Đối với phiên bản cũ không hoạt động tốt, nhưng đơn giản hơn, xem bên dưới.)

@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)]

Quá trình này xử lý danh sách vật phẩm mười triệu trong 75ms, gần gấp 50 lần so với phiên bản dựa trên danh sách được viết bằng Python thuần túy.

Đối với phiên bản chậm hơn nhưng có thể dễ đọc hơn, đây là những gì tôi đã có trước đây, dựa trên hỗ trợ thử nghiệm được thêm gần đây cho "danh sách đã nhập" có kích thước động, cho phép chúng tôi điền vào mỗi thùng theo cách không theo thứ tự nhanh hơn nhiều.

Cuộc vật lộn này với numbađộng cơ suy luận kiểu một chút, và tôi chắc chắn có cách tốt hơn để xử lý phần đó. Điều này cũng hóa ra là chậm hơn gần 10 lần so với ở trên.

@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

Tôi đã thử nghiệm những điều sau đây:

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()

Tôi cũng đã thử nghiệm chúng với phiên bản cython được biên dịch trước tương tự enum_bins_numba_buffer(được mô tả chi tiết bên dưới).

Trong danh sách mười triệu ints ngẫu nhiên ( ints = np.random.randint(0, 100, 10000000)) tôi nhận được các kết quả sau:

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)

Thật ấn tượng, cách làm việc này với numbahiệu suất vượt trội so với một cythonphiên bản của cùng chức năng, ngay cả khi đã kiểm tra giới hạn. Tôi chưa đủ quen thuộc pythranđể thử nghiệm phương pháp này bằng cách sử dụng nó, nhưng tôi rất muốn xem một so sánh. Có vẻ như dựa trên sự tăng tốc này, pythranphiên bản cũng có thể nhanh hơn một chút với phương pháp này.

Đây là cythonphiên bản để tham khảo, với một số hướng dẫn xây dựng. Khi bạn đã cythoncài đặt, bạn sẽ cần một setup.pytệp đơn giản như thế này:

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()]
)

Và mô-đun 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)]

Với hai tệp này trong thư mục làm việc của bạn, hãy chạy lệnh này:

python setup.py build_ext --inplace

Sau đó bạn có thể nhập hàm bằng cách sử dụng from enum_bins_cython import enum_bins_cython.

Đây là một cách thực sự kỳ lạ để làm điều này thật tồi tệ, nhưng tôi thấy thật buồn cười khi không chia sẻ - và tất cả 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]]

EDIT: đây là phương pháp tốt nhất tôi có thể tìm thấy dọc theo con đường này. Nó vẫn chậm hơn 10 lần so với argsortgiải pháp của @PaulPanzer :

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)

Bạn có thể làm điều đó bằng cách tạo một từ điển các số, các khóa sẽ là các số và giá trị sẽ là các chỉ số mà số đó nhìn thấy, đây là một trong những cách nhanh nhất để làm điều đó, bạn có thể thấy mã dưới đây:

>>> 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] 

Mã giả:

1. lấy "số mảng 1d trong mảng 2d", bằng cách trừ giá trị tối thiểu của mảng numpy của bạn khỏi giá trị tối đa và sau đó cộng một. Trong trường hợp của bạn, nó sẽ là 5-0 + 1 = 6

2. khởi tạo một mảng 2d với số lượng mảng 1d bên trong nó. Trong trường hợp của bạn, khởi tạo một mảng 2d với 6 mảng 1d trong đó. Mỗi mảng 1d tương ứng với một phần tử duy nhất trong mảng numpy của bạn, ví dụ, mảng 1d đầu tiên sẽ tương ứng với '0', mảng 1d thứ hai sẽ tương ứng với '1', ...

3. lặp qua mảng numpy của bạn, đặt chỉ mục của phần tử vào mảng 1d tương ứng bên phải. Trong trường hợp của bạn, chỉ mục của phần tử đầu tiên trong mảng numpy của bạn sẽ được đưa vào mảng 1d thứ hai, chỉ mục của phần tử thứ hai trong mảng numpy của bạn sẽ được đưa vào mảng 1d thứ ba, ....

Mã giả này sẽ mất thời gian tuyến tính để chạy vì nó phụ thuộc vào độ dài của mảng numpy của bạn.

Điều này cung cấp cho bạn chính xác những gì bạn muốn và sẽ mất khoảng 2,5 giây cho 10.000.000 trên máy của tôi:

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()))

Vì vậy, đưa ra một danh sách các yếu tố, bạn muốn tạo các cặp (yếu tố, chỉ mục). Trong thời gian tuyến tính, điều này có thể được thực hiện như sau:

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())

Điều này sẽ mất thời gian O (n). Tôi không thể nghĩ ra một giải pháp nhanh hơn như bây giờ, nhưng sẽ cập nhật ở đây nếu tôi làm.