Итак, предположим, что у меня есть этот 2D-массив в python
a = [[1,2]
[2,3]
[3,2]
[1,3]]
Как получить все записи массива с тем же значением строки и сохранить их в новой матрице. Например, у меня будет
b = [1,2]
[1,3]
после запроса.
Мой подход - это b = [a[i] for я in a if a[i][0] == 1][0]]
но он, похоже, не работает?
Я новичок в Python, и вся информация о разрезе индексов смущает. Спасибо!
Поскольку вы отметили 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]
.
Синтаксис for я 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 я in a if a[i][0]==1]
пытается индексировать элемент a с каждым из элементов a. Но поскольку каждый элемент 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
) и т.д. И т.д.
b = [x for x in a if x[0] == 1]