Получить все строки с одинаковыми значениями в Python?

Question

Итак, предположим, у меня есть этот двумерный массив в python

a = [[1,2]
     [2,3]
     [3,2]
     [1,3]]

Как получить все записи массива с одинаковым значением строки и сохранить их в новой матрице.Например, у меня будет

b = [1,2]
    [1,3]

после запроса.

Мой подход b = [a[i] for i in a if a[i][0] == 1][0]] , но, похоже, он не работает?

Я новичок в Python, и вся эта нарезка индекса немного сбивает с толку.Спасибо!

Answer 1

Синтаксис for i in a дает вам фактические элементы в списке ... например, так:

list_of_strs = ['first', 'second', 'third']
first_letters = [s[0] for s in list_of_strs]
# first_letters == ['f', 's', 't']

То, что вы на самом деле делаете с b = [a[i] for i in a if a[i][0]==1], пытается индексировать элемент с каждымэлементов.Но поскольку каждый элемент a сам по себе является списком, это не будет работать (вы не можете индексировать списки с другими списками)

Что-то вроде этого должно работать:

b = [row for row in a if row[0] == 1]

Бонусные баллы, если вы пишете это как функцию, так что вы можете выбрать, что именно вы хотите фильтровать.

Если вы много работаете с массивами, вы также можете проверить библиотеку numpy .С помощью numpy вы можете делать такие вещи:

import numpy as np
a = np.array([[1,2], [2,3], [3,2], [1,3]])
b = a[a[:,0] == 1]

Последняя строка в основном индексирует исходный массив a с логическим массивом, определенным в первом наборе квадратных скобок.Он очень гибкий, поэтому вы также можете изменить его, чтобы отфильтровать второй элемент, отфильтровать другие условия (например, > some_number) и т. Д. И т. Д.

Answer 2

Поскольку вы пометили numpy , вы можете выполнить эту задачу с массивами NumPy.Сначала определите ваш массив:

a = np.array([[1, 2],
              [2, 3],
              [3, 2],
              [1, 3]])

Для всех уникальных значений в первом столбце вы можете использовать словарь.Это полезно во избежание дублирования операций.

d = {i: a[a[:, 0] == i] for i in np.unique(a[:, 0])}

{1: array([[1, 2],
           [1, 3]]),
 2: array([[2, 3]]),
 3: array([[3, 2]])}

Затем получите доступ к вашему массиву, где первый столбец равен 1, с помощью d[1].

. Для одного запроса вы можете просто использовать a[a[:, 0] == 1].