ฉันจะรับดัชนีค่าสูงสุด N ในอาร์เรย์ NumPy ได้อย่างไร


482

NumPy np.argmaxเสนอวิธีที่จะได้ดัชนีของค่าสูงสุดของอาร์เรย์ผ่าน

ฉันต้องการสิ่งที่คล้ายกัน แต่คืนค่าดัชนีของNค่าสูงสุด

ตัวอย่างเช่นถ้าฉันมีอาร์เรย์[1, 3, 2, 4, 5], function(array, n=3)จะกลับมาดัชนีซึ่งสอดคล้องกับองค์ประกอบ[4, 3, 1][5, 4, 3]



4
คำถามของคุณยังไม่ชัดเจน ตัวอย่างเช่นสิ่งที่จะดัชนี (ที่คุณคาดหวัง) ที่จะให้array([5, 1, 5, 5, 2, 3, 2, 4, 1, 5]), เล็กน้อยn= 3? ซึ่งหนึ่งในทุกทางเลือกเช่น[0, 2, 3], [0, 2, 9], ...จะเป็นหนึ่งที่ถูกต้อง? โปรดอธิบายรายละเอียดเพิ่มเติมเกี่ยวกับข้อกำหนดเฉพาะของคุณ ขอบคุณ
กิน

@ เรียบร้อยฉันไม่สนใจจริงๆว่าควรส่งคืนรายการใดในกรณีนี้โดยเฉพาะ แม้ว่ามันจะดูสมเหตุสมผลที่จะส่งคืนสิ่งแรกที่พบนั่นไม่ใช่ข้อกำหนดสำหรับฉัน
Alexis Métaireau

argsortอาจเป็นทางเลือกที่ทำงานได้หากคุณไม่สนใจคำสั่งของดัชนีที่ส่งคืน ดูคำตอบของฉันด้านล่าง
สีฟ้า

คำตอบ:


347

สิ่งที่ง่ายที่สุดที่ฉันสามารถทำได้คือ:

In [1]: import numpy as np

In [2]: arr = np.array([1, 3, 2, 4, 5])

In [3]: arr.argsort()[-3:][::-1]
Out[3]: array([4, 3, 1])

สิ่งนี้เกี่ยวข้องกับอาร์เรย์ที่สมบูรณ์ ฉันสงสัยว่าnumpyมีวิธีการเรียงลำดับบางส่วนในตัวหรือไม่ จนถึงตอนนี้ฉันไม่สามารถหามันได้

หากการแก้ปัญหานี้จะเปิดออกจะช้าเกินไป (โดยเฉพาะอย่างยิ่งสำหรับธุรกิจขนาดเล็กn) ก็อาจจะคุ้มค่าดูที่การเข้ารหัสบางสิ่งบางอย่างขึ้นมาในCython


1
บรรทัดที่ 3 สามารถเขียนได้เทียบเท่ากันarr.argsort()[-1:-4:-1]ไหม ฉันลองใช้เป็นล่ามและมันก็เกิดผลเหมือนกัน แต่ฉันก็สงสัยว่ามันจะไม่แตกหรือเปล่า
abroekhof

44
@abroekhof ใช่ว่าควรจะเทียบเท่ากับรายการหรืออาร์เรย์ใด ๆ อีกวิธีหนึ่งสามารถทำได้โดยไม่ต้องกลับรายการโดยใช้np.argsort(-arr)[:3]ซึ่งฉันพบว่าอ่านง่ายขึ้นและตรงประเด็นมากขึ้น
askewchan

6
[:: - 1] หมายถึงอะไร @NPE
1a1a11a

@ 1a1a11a หมายถึงการย้อนกลับอาร์เรย์ (ตามตัวอักษรนำสำเนาของอาร์เรย์จากนาทีที่ไม่ จำกัด ไปยังสูงสุดที่ไม่ จำกัด ในลำดับที่กลับด้าน)
FizBack

15
arr.argsort()[::-1][:n]จะดีกว่าเพราะจะส่งคืนค่าว่างn=0แทนชุดเต็ม
abora

599

รุ่น NumPy ที่ใหม่กว่า (1.8 ขึ้นไป) มีฟังก์ชันที่เรียกใช้argpartitionสำหรับสิ่งนี้ เพื่อให้ได้ดัชนีขององค์ประกอบที่ใหญ่ที่สุดสี่อย่างให้ทำ

>>> a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> a
array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> ind = np.argpartition(a, -4)[-4:]
>>> ind
array([1, 5, 8, 0])
>>> a[ind]
array([4, 9, 6, 9])

ซึ่งแตกต่างจากargsortฟังก์ชั่นนี้วิ่งในเส้นเวลาในกรณีที่เลวร้ายที่สุด a[ind]แต่ดัชนีกลับมาจะไม่ได้เรียงที่สามารถมองเห็นได้จากผลของการประเมิน หากคุณต้องการสิ่งนั้นให้เรียงลำดับภายหลัง:

>>> ind[np.argsort(a[ind])]
array([1, 8, 5, 0])

ที่จะได้รับบนสุดkองค์ประกอบในเรียงลำดับในลักษณะนี้จะใช้เวลา O ( n + kบันทึกk ) เวลา


27
@varela argpartitionทำงานในเวลาเชิงเส้น O (n) โดยใช้อัลกอริทึมintroselect การเรียงลำดับที่ตามมาจะจัดการองค์ประกอบ k เท่านั้นดังนั้นการทำงานใน O (k log k)
Fred Foo

2
หากใครสงสัยว่าวิธีการnp.argpartitionและอัลกอริทึมของน้องสาวnp.partitionนั้นทำงานอย่างไรมีคำอธิบายโดยละเอียดเพิ่มเติมในคำถามที่เชื่อมโยง: stackoverflow.com/questions/10337533/…
Ramon Martinez

7
@FredFoo: ทำไมคุณถึงใช้ -4 คุณทำอย่างนั้นเพื่อเริ่มถอยหลังหรือไม่ (เนื่องจาก k เป็นงานบวกหรือลบเหมือนกันสำหรับฉัน! มันพิมพ์เฉพาะตัวเลขที่เล็กที่สุดก่อน!
Rika

2
ใช้ @LKT a=np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])เนื่องจากรายการไพ ธ อนปกติไม่รองรับการสร้างดัชนีตามรายการซึ่งแตกต่างจากnp.array
Marawan Okasha

2
@Umangsinghal np.argpartitionรับaxisอาร์กิวเมนต์ตัวเลือก เพื่อหาดัชนีของค่า n สูงสุดของแต่ละแถว:np.argpartition(a, -n, axis=1)[-n:]
Ralph

48

เรียบง่ายขึ้น:

idx = (-arr).argsort()[:n]

โดยที่nคือจำนวนค่าสูงสุด


7
สิ่งนี้สามารถทำได้สำหรับอาร์เรย์ 2d หรือไม่? ถ้าไม่คุณอาจรู้ได้อย่างไร
Andrew Hundt

2
@AndrewHundt: เพียงใช้ (-arr) .argsort (axis = -1) [:,: n]
MiniQuark

2
ที่คล้ายกันจะเป็นarr[arr.argsort()[-n:]]แทนการปฏิเสธอาร์เรย์เพียงแค่ใช้ชิ้นส่วนขององค์ประกอบสุดท้าย n
loganjones16

35

ใช้:

>>> import heapq
>>> import numpy
>>> a = numpy.array([1, 3, 2, 4, 5])
>>> heapq.nlargest(3, range(len(a)), a.take)
[4, 3, 1]

สำหรับรายการ Python ปกติ:

>>> a = [1, 3, 2, 4, 5]
>>> heapq.nlargest(3, range(len(a)), a.__getitem__)
[4, 3, 1]

ถ้าคุณใช้งูหลาม 2 ใช้แทนxrangerange

ที่มา: heapq - อัลกอริทึมคิว Heap


2
heapq.nlargest(3, xrange(len(a)), a.take)ไม่จำเป็นต้องมีของวงที่ทุกคนที่นี่ไม่มีที่: สำหรับรายการหลามเราสามารถใช้แทน.__getitem__ .take
Ashwini Chaudhary

สำหรับอาร์เรย์ n มิติทั่วไป:A heapq.nlargest(3, range(len(A.ravel())), A.ravel().take)(ฉันหวังว่านี่จะใช้งานได้กับการดูเท่านั้นโปรดดู ( ravel vs flatten] ( stackoverflow.com/a/28930580/603003 ))
ComFreek

31

หากคุณทำงานกับอาร์เรย์หลายมิติคุณจะต้องทำให้ดัชนีแบนและคลี่คลาย:

def largest_indices(ary, n):
    """Returns the n largest indices from a numpy array."""
    flat = ary.flatten()
    indices = np.argpartition(flat, -n)[-n:]
    indices = indices[np.argsort(-flat[indices])]
    return np.unravel_index(indices, ary.shape)

ตัวอย่างเช่น:

>>> xs = np.sin(np.arange(9)).reshape((3, 3))
>>> xs
array([[ 0.        ,  0.84147098,  0.90929743],
       [ 0.14112001, -0.7568025 , -0.95892427],
       [-0.2794155 ,  0.6569866 ,  0.98935825]])
>>> largest_indices(xs, 3)
(array([2, 0, 0]), array([2, 2, 1]))
>>> xs[largest_indices(xs, 3)]
array([ 0.98935825,  0.90929743,  0.84147098])

9

ถ้าคุณไม่สนใจเกี่ยวกับการสั่งซื้อของ K-TH องค์ประกอบที่ใหญ่ที่สุดที่คุณสามารถใช้ซึ่งควรจะมีประสิทธิภาพสูงกว่าการจัดเรียงเต็มรูปแบบผ่านargpartitionargsort

K = 4 # We want the indices of the four largest values
a = np.array([0, 8, 0, 4, 5, 8, 8, 0, 4, 2])
np.argpartition(a,-K)[-K:]
array([4, 1, 5, 6])

เครดิตไปที่คำถามนี้

ฉันทำการทดสอบสองสามครั้งและดูเหมือนว่าจะมีargpartitionประสิทธิภาพสูงกว่าargsortเมื่อขนาดของอาร์เรย์และมูลค่าของ K เพิ่มขึ้น


7

สำหรับอาร์เรย์หลายมิติคุณสามารถใช้axisคำหลักเพื่อใช้การแบ่งพาร์ติชันตามแกนที่คาดไว้

# For a 2D array
indices = np.argpartition(arr, -N, axis=1)[:, -N:]

และสำหรับการคว้ารายการ:

x = arr.shape[0]
arr[np.repeat(np.arange(x), N), indices.ravel()].reshape(x, N)

แต่โปรดทราบว่าสิ่งนี้จะไม่ส่งคืนผลลัพธ์ที่เรียง ในกรณีนั้นคุณสามารถใช้np.argsort()ตามแกนที่ต้องการได้:

indices = np.argsort(arr, axis=1)[:, -N:]

# Result
x = arr.shape[0]
arr[np.repeat(np.arange(x), N), indices.ravel()].reshape(x, N)

นี่คือตัวอย่าง:

In [42]: a = np.random.randint(0, 20, (10, 10))

In [44]: a
Out[44]:
array([[ 7, 11, 12,  0,  2,  3,  4, 10,  6, 10],
       [16, 16,  4,  3, 18,  5, 10,  4, 14,  9],
       [ 2,  9, 15, 12, 18,  3, 13, 11,  5, 10],
       [14,  0,  9, 11,  1,  4,  9, 19, 18, 12],
       [ 0, 10,  5, 15,  9, 18,  5,  2, 16, 19],
       [14, 19,  3, 11, 13, 11, 13, 11,  1, 14],
       [ 7, 15, 18,  6,  5, 13,  1,  7,  9, 19],
       [11, 17, 11, 16, 14,  3, 16,  1, 12, 19],
       [ 2,  4, 14,  8,  6,  9, 14,  9,  1,  5],
       [ 1, 10, 15,  0,  1,  9, 18,  2,  2, 12]])

In [45]: np.argpartition(a, np.argmin(a, axis=0))[:, 1:] # 1 is because the first item is the minimum one.
Out[45]:
array([[4, 5, 6, 8, 0, 7, 9, 1, 2],
       [2, 7, 5, 9, 6, 8, 1, 0, 4],
       [5, 8, 1, 9, 7, 3, 6, 2, 4],
       [4, 5, 2, 6, 3, 9, 0, 8, 7],
       [7, 2, 6, 4, 1, 3, 8, 5, 9],
       [2, 3, 5, 7, 6, 4, 0, 9, 1],
       [4, 3, 0, 7, 8, 5, 1, 2, 9],
       [5, 2, 0, 8, 4, 6, 3, 1, 9],
       [0, 1, 9, 4, 3, 7, 5, 2, 6],
       [0, 4, 7, 8, 5, 1, 9, 2, 6]])

In [46]: np.argpartition(a, np.argmin(a, axis=0))[:, -3:]
Out[46]:
array([[9, 1, 2],
       [1, 0, 4],
       [6, 2, 4],
       [0, 8, 7],
       [8, 5, 9],
       [0, 9, 1],
       [1, 2, 9],
       [3, 1, 9],
       [5, 2, 6],
       [9, 2, 6]])

In [89]: a[np.repeat(np.arange(x), 3), ind.ravel()].reshape(x, 3)
Out[89]:
array([[10, 11, 12],
       [16, 16, 18],
       [13, 15, 18],
       [14, 18, 19],
       [16, 18, 19],
       [14, 14, 19],
       [15, 18, 19],
       [16, 17, 19],
       [ 9, 14, 14],
       [12, 15, 18]])

ฉันคิดว่าคุณสามารถลดความซับซ้อนของการจัดทำดัชนีที่นี่โดยใช้np.take_along_axis(ซึ่งมีแนวโน้มที่ไม่ได้อยู่เมื่อคุณตอบคำถามนี้)
เอริค

4

การดำเนินการนี้จะเร็วกว่าการจัดเรียงแบบเต็มรูปแบบขึ้นอยู่กับขนาดของอาเรย์ดั้งเดิมและขนาดที่คุณเลือก:

>>> A = np.random.randint(0,10,10)
>>> A
array([5, 1, 5, 5, 2, 3, 2, 4, 1, 0])
>>> B = np.zeros(3, int)
>>> for i in xrange(3):
...     idx = np.argmax(A)
...     B[i]=idx; A[idx]=0 #something smaller than A.min()
...     
>>> B
array([0, 2, 3])

แน่นอนว่าเกี่ยวข้องกับการดัดแปลงอาร์เรย์เดิมของคุณ ซึ่งคุณสามารถแก้ไข (ถ้าจำเป็น) โดยการทำสำเนาหรือเปลี่ยนกลับค่าเดิม ... แล้วแต่ว่าราคาใดจะถูกกว่าสำหรับเคสที่คุณใช้


FWIW โซลูชันของคุณจะไม่ให้โซลูชันที่ชัดเจนในทุกสถานการณ์ OP ควรอธิบายวิธีจัดการกรณีที่ไม่คลุมเครือเหล่านี้ ขอบคุณ
กิน

@eat คำถามของ OP ค่อนข้างคลุมเครือ อย่างไรก็ตามการใช้งานนั้นไม่ได้เปิดกว้างสำหรับการตีความ :) OP ควรอ้างถึงคำจำกัดความของ np.argmax docs.scipy.org/doc/numpy/reference/generated/numpy.argmax.htmlเพื่อให้แน่ใจว่าโซลูชันเฉพาะนี้ตรงตามข้อกำหนด เป็นไปได้ว่าวิธีการแก้ปัญหาใด ๆ ก็ตามที่เป็นไปตามข้อกำหนดของ OP นั้นเป็นที่ยอมรับได้
พอล

ก็อาจพิจารณาการใช้งานargmax(.)ที่ไม่คลุมเครือเช่นกัน (IMHO จะพยายามติดตามตรรกะการลัดวงจรบางประเภท แต่น่าเสียดายที่ไม่สามารถให้พฤติกรรมที่ยอมรับได้ในระดับสากล) ขอบคุณ
กิน

3

เมธอดnp.argpartitionจะส่งคืนดัชนีที่ใหญ่ที่สุด k ทำการเรียงลำดับแบบโลคัลและเร็วกว่าnp.argsort(ดำเนินการเรียงแบบเต็ม) เมื่ออาร์เรย์มีขนาดใหญ่มาก แต่ดัชนีกลับมาจะไม่ได้อยู่ในจากน้อยไปมาก / ลำดับถัดลงมา สมมติด้วยตัวอย่าง:

ป้อนคำอธิบายภาพที่นี่

เราจะเห็นว่าถ้าคุณต้องการดัชนีการเรียงลำดับบน k ที่เข้มงวดมากขึ้นnp.argpartitionจะไม่ส่งคืนสิ่งที่คุณต้องการ

นอกเหนือจากการเรียงลำดับด้วยตนเองหลังจาก np.argpartition โซลูชันของฉันคือการใช้ PyTorch torch.topkซึ่งเป็นเครื่องมือสำหรับการสร้างเครือข่ายประสาทเทียมให้ API แบบ NumPy ที่เหมือนกันทั้ง CPU และ GPU มันเร็วเท่ากับ NumPy กับ MKL และให้การสนับสนุน GPU หากคุณต้องการการคำนวณเมทริกซ์ / เวกเตอร์ขนาดใหญ่

รหัสดัชนี k บนสุดขึ้น / ลงอย่างเข้มงวดจะเป็น:

ป้อนคำอธิบายภาพที่นี่

โปรดทราบว่าtorch.topkยอมรับเมตริกซ์ไฟฉายและผลตอบแทนทั้งค่า k ด้านบนและดัชนี k torch.Tensorชั้นนำในประเภท เช่นเดียวกับ np, torch.topk ยังยอมรับอาร์กิวเมนต์แกนเพื่อให้คุณสามารถจัดการอาร์เรย์ / เทนเซอร์หลายมิติ


2

ใช้:

from operator import itemgetter
from heapq import nlargest
result = nlargest(N, enumerate(your_list), itemgetter(1))

ตอนนี้resultรายการจะมีn tuples ( index, value) ที่valueถูกขยายให้ใหญ่สุด


2

ใช้:

def max_indices(arr, k):
    '''
    Returns the indices of the k first largest elements of arr
    (in descending order in values)
    '''
    assert k <= arr.size, 'k should be smaller or equal to the array size'
    arr_ = arr.astype(float)  # make a copy of arr
    max_idxs = []
    for _ in range(k):
        max_element = np.max(arr_)
        if np.isinf(max_element):
            break
        else:
            idx = np.where(arr_ == max_element)
        max_idxs.append(idx)
        arr_[idx] = -np.inf
    return max_idxs

มันยังทำงานร่วมกับอาร์เรย์ 2 มิติ ตัวอย่างเช่น,

In [0]: A = np.array([[ 0.51845014,  0.72528114],
                     [ 0.88421561,  0.18798661],
                     [ 0.89832036,  0.19448609],
                     [ 0.89832036,  0.19448609]])
In [1]: max_indices(A, 8)
Out[1]:
    [(array([2, 3], dtype=int64), array([0, 0], dtype=int64)),
     (array([1], dtype=int64), array([0], dtype=int64)),
     (array([0], dtype=int64), array([1], dtype=int64)),
     (array([0], dtype=int64), array([0], dtype=int64)),
     (array([2, 3], dtype=int64), array([1, 1], dtype=int64)),
     (array([1], dtype=int64), array([1], dtype=int64))]

In [2]: A[max_indices(A, 8)[0]][0]
Out[2]: array([ 0.89832036])

ใช้งานได้ดี แต่ให้ผลลัพธ์มากกว่านี้ถ้าคุณมีค่าซ้ำ (สูงสุด) ในอาร์เรย์ A ของคุณฉันคาดหวังผลลัพธ์ k อย่างแน่นอน แต่ในกรณีที่มีค่าซ้ำกันคุณจะได้ผลลัพธ์มากกว่า k
กุยโด้

ฉันแก้ไขรหัสเล็กน้อย รายการดัชนีที่ส่งคืนมีความยาวเท่ากับ k หากคุณมีรายการซ้ำพวกเขาจะถูกจัดกลุ่มเป็น tuple เดียว
X Æ A-12

1

bottleneck มีฟังก์ชั่นการเรียงลำดับบางส่วนถ้าค่าใช้จ่ายในการเรียงลำดับอาร์เรย์ทั้งหมดเพียงเพื่อให้ได้ค่า N ที่มากที่สุดนั้นมากเกินไป

ฉันไม่รู้อะไรเกี่ยวกับโมดูลนี้ ฉันเพียงแค่ numpy partial sortgoogled


ฉันพบว่าไม่มีฟังก์ชั่นการเรียงลำดับบางส่วนในคอขวดมีฟังก์ชั่นพาร์ทิชัน แต่นี่ไม่ได้เรียงลำดับ
nbecker

1

ต่อไปนี้เป็นวิธีที่ง่ายมากในการดูองค์ประกอบสูงสุดและตำแหน่งของมัน นี่axisคือโดเมน axis= 0 หมายถึงจำนวนสูงสุดของคอลัมน์ที่ชาญฉลาดและaxis= 1 หมายถึงจำนวนสูงสุดของแถวที่ชาญฉลาดสำหรับกรณี 2D และสำหรับมิติที่สูงขึ้นมันขึ้นอยู่กับคุณ

M = np.random.random((3, 4))
print(M)
print(M.max(axis=1), M.argmax(axis=1))

ผมใช้การเชื่อมโยงนี้jakevdp.github.io/PythonDataScienceHandbook/...
เสรีนิยม

0

np.uniqueผมพบว่ามันใช้งานง่ายที่สุดในการใช้

ความคิดคือว่าวิธีการที่ไม่ซ้ำกันส่งกลับดัชนีของค่าที่ป้อนเข้า จากค่าสูงสุดที่ไม่ซ้ำกันและบ่งชี้ตำแหน่งของค่าเดิมสามารถสร้างใหม่ได้

multi_max = [1,1,2,2,4,0,0,4]
uniques, idx = np.unique(multi_max, return_inverse=True)
print np.squeeze(np.argwhere(idx == np.argmax(uniques)))
>> [4 7]

0

ฉันคิดว่าวิธีที่มีประสิทธิภาพเวลามากที่สุดคือวนซ้ำแบบแมนนวลผ่านอาเรย์และรักษา k-size min-heap อย่างที่คนอื่น ๆ พูดถึง

และฉันยังมาพร้อมกับวิธีการดุร้ายกำลัง:

top_k_index_list = [ ]
for i in range(k):
    top_k_index_list.append(np.argmax(my_array))
    my_array[top_k_index_list[-1]] = -float('inf')

ตั้งค่าองค์ประกอบที่ใหญ่ที่สุดเป็นค่าลบขนาดใหญ่หลังจากที่คุณใช้ argmax เพื่อรับดัชนี จากนั้นการเรียก argmax ครั้งต่อไปจะคืนค่าองค์ประกอบที่ใหญ่เป็นอันดับสอง และคุณสามารถบันทึกค่าดั้งเดิมขององค์ประกอบเหล่านี้และกู้คืนได้หากคุณต้องการ


0

รหัสนี้ใช้ได้กับอาร์เรย์เมทริกซ์ที่มีจำนวนมาก:

mat = np.array([[1, 3], [2, 5]]) # numpy matrix

n = 2  # n
n_largest_mat = np.sort(mat, axis=None)[-n:] # n_largest 
tf_n_largest = np.zeros((2,2), dtype=bool) # all false matrix
for x in n_largest_mat: 
  tf_n_largest = (tf_n_largest) | (mat == x) # true-false  

n_largest_elems = mat[tf_n_largest] # true-false indexing 

สิ่งนี้สร้างการจัดทำดัชนีเมทริกซ์ n_largest ที่เป็นเท็จซึ่งยังสามารถแยกองค์ประกอบ n_largest จากอาเรย์เมทริกซ์

โดยการใช้ไซต์ของเรา หมายความว่าคุณได้อ่านและทำความเข้าใจนโยบายคุกกี้และนโยบายความเป็นส่วนตัวของเราแล้ว
Licensed under cc by-sa 3.0 with attribution required.