Chọn các hàng từ một số NumPy

Question

Tôi muốn chỉ chọn một số hàng nhất định từ một mảng NumPy dựa trên giá trị trong cột thứ hai. Ví dụ, mảng kiểm tra này có số nguyên từ 1 đến 10 trong cột thứ hai.

>>> test = numpy.array([numpy.arange(100), numpy.random.randint(1, 11, 100)]).transpose() 
>>> test[:10, :] 
array([[ 0, 6], 
     [ 1, 7], 
     [ 2, 10], 
     [ 3, 4], 
     [ 4, 1], 
     [ 5, 10], 
     [ 6, 6], 
     [ 7, 4], 
     [ 8, 6], 
     [ 9, 7]]) 

Nếu tôi muốn chỉ hàng trong đó giá trị thứ hai là 4, nó rất dễ dàng:

>>> test[test[:, 1] == 4] 
array([[ 3, 4], 
     [ 7, 4], 
     [16, 4], 
     ... 
     [81, 4], 
     [83, 4], 
     [88, 4]]) 

Nhưng làm thế nào để đạt được kết quả tương tự khi có nhiều hơn một giá trị mong muốn?

Danh sách mong muốn có thể có độ dài tùy ý. Ví dụ, tôi có thể muốn tất cả các hàng nơi cột số thứ hai là 2, 4 hoặc 6:

>>> wanted = [2, 4, 6] 

Cách duy nhất mà tôi đã đưa ra là sử dụng danh sách hiểu và sau đó chuyển đổi này trở lại vào một mảng và dường như quá phức tạp, mặc dù nó hoạt động:

>>> test[numpy.array([test[x, 1] in wanted for x in range(len(test))])] 
array([[ 0, 6], 
     [ 3, 4], 
     [ 6, 6], 
     ... 
     [90, 2], 
     [91, 6], 
     [92, 2]]) 

Có cách nào tốt hơn để làm điều này trong chính NumPy mà tôi bị thiếu không?

Answer 1

test[numpy.logical_or.reduce([test[:,1] == x for x in wanted])] 

Kết quả phải nhanh hơn phiên bản gốc vì NumPy thực hiện vòng lặp bên trong thay cho Python.

Answer 2

Các giải pháp sau đây cần được nhanh hơn so với giải pháp Amnon như wanted được lớn hơn:

wanted_set = set(wanted) # Much faster look up than with lists, for larger lists 

@numpy.vectorize 
def selected(elmt): return elmt in wanted_set # Or: selected = numpy.vectorize(wanted_set.__contains__) 

print test[selected(test[:, 1])] 

Trong thực tế, nó có lợi thế là tìm kiếm thông qua các mảng test chỉ lần (thay vì len(wanted) lần). Nó cũng sử dụng phần tử nhanh có sẵn của Python tra cứu trong bộ, nhanh hơn nhiều so với danh sách này. Nó cũng nhanh vì nó sử dụng các vòng lặp nhanh của Numpy. Bạn cũng nhận được tối ưu hóa của toán tử in: khi kết hợp phần tử wanted, các yếu tố còn lại không phải được kiểm tra (trái ngược với cách tiếp cận "hợp lý hoặc" của Amnon, tất cả các phần tử trong wanted đều được kiểm tra bất kể) .

Ngoài ra, bạn có thể sử dụng sau một lót, mà cũng đi qua mảng của bạn chỉ một lần:

test[numpy.apply_along_axis(lambda x: x[1] in wanted, 1, test)] 

Đây là nhiều chậm hơn nhiều, tuy nhiên, vì điều này chiết xuất phần tử trong cột thứ hai tại mỗi lặp lại (thay vì thực hiện nó trong một lần truyền, như trong giải pháp đầu tiên của câu trả lời này).

Answer 3

này nhanh hơn hai lần so với biến thể Amnon cho len (thử nghiệm) = 1000:

wanted = (2,4,6) 
wanted2 = numpy.expand_dims(wanted, 1) 
print test[numpy.any(test[:, 1] == wanted2, 0), :] 

Answer 4

numpy.in1d là những gì bạn đang tìm kiếm:

print test[numpy.in1d(test[:,1], wanted)] 

Nó dễ dàng nên có giải pháp nhanh nhất nếu muốn lớn; cộng với, nó là một trong những dễ đọc nhất, id nói.