Метод compress
Метод compress класса ndarray возвращает новый массив,
состоящий из элементов исходного массива, для которых соответствующие
значения в массиве условий истинны. Первым параметром метод принимает
массив условий condition, который должен быть одномерным
или иметь форму, допускающую трансляцию. Вторым параметром можно
задать ось axis, вдоль которой выполняется сжатие.
Синтаксис
arr.compress(condition, axis=None, out=None)
Пример
Давайте выберем элементы одномерного массива, которые больше трёх:
import numpy as np
arr = np.array([1, 2, 3, 4, 5])
cond = arr > 3
res = arr.compress(cond)
print(res)
Результат выполнения кода:
[4 5]
Пример
Вторым параметром метод принимает ось, вдоль которой выполняется сжатие. По умолчанию массив условий применяется к сплющенному массиву. Давайте выберем строки двумерного массива, в которых сумма элементов больше десяти:
import numpy as np
arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
cond = arr.sum(axis=1) > 10
res = arr.compress(cond, axis=0)
print(res)
Результат выполнения кода:
[[4 5 6]
[7 8 9]]
Пример
Метод также можно использовать с массивом условий, который имеет ту же длину, что и ось, вдоль которой происходит сжатие. Давайте выберем столбцы, в которых все элементы положительные:
import numpy as np
arr = np.array([[-1, 2, 3], [-4, 5, 6], [7, 8, 9]])
cond = np.all(arr > 0, axis=0)
res = arr.compress(cond, axis=1)
print(res)
Результат выполнения кода:
[[2 3]
[5 6]
[8 9]]
Пример
Если массив условий одномерный и его длина совпадает с размером соответствующей оси, то можно использовать булев массив для выбора:
import numpy as np
arr = np.array([10, 20, 30, 40, 50])
cond = np.array([True, False, True, False, True])
res = arr.compress(cond)
print(res)
Результат выполнения кода:
[10 30 50]