Numpy: Как найти индексы нескольких значений в массиве?

Краткое описание NumPy и его массивов

NumPy (Numerical Python) является фундаментальной библиотекой для научных вычислений в Python. Ее центральным объектом является многомерный массив ndarray, который обеспечивает эффективное хранение и манипулирование большими наборами числовых данных. Массивы NumPy оптимизированы для быстрых операций, включая векторные операции, что делает их незаменимыми в анализе данных, машинном обучении и других областях, требующих интенсивных вычислений.

Задача поиска индексов нескольких определенных значений

Одной из частых задач при работе с массивами NumPy является поиск позиций (индексов) элементов, которые соответствуют определенным критериям. В то время как поиск индекса одного конкретного значения или элементов, удовлетворяющих одному простому условию (например, больше определенного числа), достаточно тривиален, задача усложняется, когда необходимо найти индексы элементов, которые принадлежат набору из нескольких предопределенных значений или удовлетворяют комбинации различных условий.

Обзор различных подходов к решению задачи

Для решения этой задачи в NumPy существует несколько эффективных методов. Наиболее распространенные и производительные подходы включают использование функций np.where и np.isin. Каждый из этих методов имеет свои особенности и оптимальные сценарии применения, а их комбинация позволяет решать достаточно сложные задачи фильтрации и поиска индексов.

Использование np.where для поиска индексов

Базовый синтаксис np.where

Функция np.where является мощным инструментом для выполнения условных операций над элементами массива. Ее базовый синтаксис выглядит как np.where(condition, x, y), где condition — это булев массив (или массив, который может быть преобразован в булев), а x и y — значения или массивы, используемые для выбора элементов в зависимости от того, истинно (x) или ложно (y) соответствующее условие.

Если x и y опущены (np.where(condition)), функция возвращает кортеж массивов индексов элементов, для которых условие истинно. Это и делает ее полезной для поиска индексов.

Поиск индексов элементов, удовлетворяющих одному условию

Простейшее применение np.where для поиска индексов — это использование булева массива, полученного сравнением, в качестве условия.

import numpy as np

def find_indices_single_condition(data_array: np.ndarray, threshold: float) -> np.ndarray:
    """
    Находит индексы элементов в массиве, которые больше заданного порога.

    Args:
        data_array: Входной NumPy массив.
        threshold: Пороговое значение.

    Returns:
        Массив индексов элементов, удовлетворяющих условию.
    """
    # Создаем булев массив, где True для элементов > threshold
    condition: np.ndarray = data_array > threshold

    # Используем np.where для получения индексов, где condition истинно
    # np.where(condition) возвращает кортеж массивов индексов
    indices_tuple: tuple[np.ndarray, ...] = np.where(condition)

    # Для одномерного массива возвращается кортеж с одним массивом
    return indices_tuple[0]

# Пример использования (допустим, это клики по рекламе)
clicks_per_day = np.array([150, 210, 80, 350, 190, 420, 110])
min_clicks_threshold = 200

high_click_days_indices = find_indices_single_condition(clicks_per_day, min_clicks_threshold)
# print(f"Дни с количеством кликов выше {min_clicks_threshold}: индексы {high_click_days_indices}")
# Вывод: Дни с количеством кликов выше 200: индексы [1 3 5]

Поиск индексов элементов, удовлетворяющих нескольким условиям (логические операции)

np.where также легко справляется с поиском индексов элементов, удовлетворяющих нескольким условиям одновременно или хотя бы одному из них. Для этого используются стандартные логические операторы NumPy: & (логическое И), | (логическое ИЛИ), ~ (логическое НЕ). Важно заключать каждое отдельное условие в скобки из-за приоритета операторов.

import numpy as np

def find_indices_multiple_conditions(data_array: np.ndarray, lower_bound: float, upper_bound: float) -> np.ndarray:
    """
    Находит индексы элементов в массиве, которые находятся в заданном диапазоне (включая границы).

    Args:
        data_array: Входной NumPy массив.
        lower_bound: Нижняя граница диапазона.
        upper_bound: Верхняя граница диапазона.

    Returns:
        Массив индексов элементов, удовлетворяющих условию.
    """
    # Создаем булев массив для элементов >= lower_bound И <= upper_bound
    condition: np.ndarray = (data_array >= lower_bound) & (data_array <= upper_bound)

    # Получаем индексы
    indices_tuple: tuple[np.ndarray, ...] = np.where(condition)
    return indices_tuple[0]

# Пример использования (допустим, время на сайте в секундах)
session_times_seconds = np.array([30, 125, 5, 60, 90, 250, 45, 180])
min_time = 60
max_time = 150

engaged_sessions_indices = find_indices_multiple_conditions(session_times_seconds, min_time, max_time)
# print(f"Сессии длительностью от {min_time} до {max_time} секунд: индексы {engaged_sessions_indices}")
# Вывод: Сессии длительностью от 60 до 150 секунд: индексы [3 4]

Примеры использования np.where с различными типами данных

np.where работает не только с числовыми данными, но и с булевыми, строковыми и другими типами данных, поддерживающими операции сравнения.

import numpy as np

def find_indices_string_condition(data_array: np.ndarray, target_string: str) -> np.ndarray:
    """
    Находит индексы элементов в массиве строк, равных заданной строке.

    Args:
        data_array: Входной NumPy массив строк.
        target_string: Целевая строка для поиска.

    Returns:
        Массив индексов элементов, равных целевой строке.
    """
    # Создаем булев массив, где True для элементов, равных target_string
    condition: np.ndarray = data_array == target_string

    # Получаем индексы
    indices_tuple: tuple[np.ndarray, ...] = np.where(condition)
    return indices_tuple[0]

# Пример использования (допустим, категории продуктов)
product_categories = np.array(['Электроника', 'Одежда', 'Книги', 'Электроника', 'Дом', 'Книги'])
target_category = 'Электроника'

electronics_indices = find_indices_string_condition(product_categories, target_category)
# print(f"Индексы продуктов категории '{target_category}': {electronics_indices}")
# Вывод: Индексы продуктов категории 'Электроника': [0 3]

Использование np.isin для поиска индексов

Описание функции np.isin и ее назначения

Функция np.isin(element, test_elements) проверяет, присутствует ли каждый элемент из первого входного массива (element) во втором входном массиве (test_elements). Она возвращает булев массив той же формы, что и element, где True указывает на то, что соответствующий элемент найден в test_elements.

Эта функция идеально подходит для задачи поиска индексов элементов, которые являются одним из нескольких предопределенных значений.

Применение np.isin для определения наличия нескольких значений в массиве

np.isin напрямую решает задачу проверки членства в наборе значений, возвращая булеву маску.

import numpy as np

def create_isin_mask(data_array: np.ndarray, target_values: list) -> np.ndarray:
    """
    Создает булеву маску для элементов массива, которые находятся в списке целевых значений.

    Args:
        data_array: Входной NumPy массив.
        target_values: Список значений для проверки членства.

    Returns:
        Булев массив той же формы, что и data_array.
    """
    # Создаем булеву маску: True, если элемент data_array находится в target_values
    isin_mask: np.ndarray = np.isin(data_array, target_values)
    return isin_mask

# Пример использования (допустим, коды ошибок веб-сервера)
status_codes = np.array([200, 404, 500, 301, 200, 403, 404, 200])
error_codes = [403, 404, 500]

error_mask = create_isin_mask(status_codes, error_codes)
# print(f"Булева маска для кодов ошибок: {error_mask}")
# Вывод: Булева маска для кодов ошибок: [False  True  True False False  True  True False]

Получение индексов элементов, найденных с помощью np.isin

Сама по себе np.isin возвращает маску. Чтобы получить индексы, нужно передать эту булеву маску в функцию np.where, как показано ранее, или использовать булеву индексацию (что по сути одно и то же с точки зрения результата поиска индексов).

Реклама
import numpy as np

def find_indices_isin(data_array: np.ndarray, target_values: list) -> np.ndarray:
    """
    Находит индексы элементов в массиве, которые находятся в списке целевых значений, используя np.isin.

    Args:
        data_array: Входной NumPy массив.
        target_values: Список значений для проверки членства.

    Returns:
        Массив индексов элементов, удовлетворяющих условию.
    """
    # Создаем булеву маску с помощью np.isin
    isin_mask: np.ndarray = np.isin(data_array, target_values)

    # Используем np.where с полученной маской для получения индексов
    indices_tuple: tuple[np.ndarray, ...] = np.where(isin_mask)
    return indices_tuple[0]

# Пример использования (продолжение примера с кодами ошибок)
status_codes = np.array([200, 404, 500, 301, 200, 403, 404, 200])
error_codes = [403, 404, 500]

error_indices = find_indices_isin(status_codes, error_codes)
# print(f"Индексы кодов ошибок {error_codes}: {error_indices}")
# Вывод: Индексы кодов ошибок [403, 404, 500]: [1 2 5 6]

Преимущества и недостатки использования np.isin

  • Преимущества: np.isin очень эффективна и читаема, когда нужно проверить принадлежность элементов набору дискретных значений. Она оптимизирована для этой конкретной задачи.
  • Недостатки: Производительность np.isin может снижаться, если список test_elements очень большой, так как требует проверки каждого элемента из element против каждого элемента из test_elements (хотя NumPy использует оптимизированные алгоритмы, такие как хеширование, для ускорения). Также она неприменима для условий сравнения (больше/меньше).

Комбинирование np.where и np.isin

Объединение возможностей np.where и np.isin для сложных условий

Часто возникает необходимость найти индексы элементов, которые удовлетворяют как условию членства в наборе (np.isin), так и другим критериям сравнения. Например, найти индексы транзакций, которые были совершены из определенного списка стран и сумма которых превышает заданное значение. В таких случаях np.where и np.isin прекрасно дополняют друг друга.

Можно сначала получить булеву маску с помощью np.isin, а затем скомбинировать ее с другими булевыми масками, полученными с помощью сравнений, используя логические операторы (&, |) внутри np.where.

Примеры реализации комбинированных запросов

Рассмотрим пример из области веб-аналитики: найти индексы пользователей, которые зашли на сайт с определенного списка браузеров и провели на сайте больше 60 секунд.

import numpy as np

def find_indices_complex_condition(browsers: np.ndarray, session_times: np.ndarray, target_browsers: list, min_time_seconds: int) -> np.ndarray:
    """
    Находит индексы сессий, совершенных с целевых браузеров И длившихся дольше заданного времени.

    Args:
        browsers: Массив строк с названиями браузеров.
        session_times: Массив чисел с длительностью сессий в секундах.
        target_browsers: Список целевых названий браузеров.
        min_time_seconds: Минимальная длительность сессии в секундах.

    Returns:
        Массив индексов сессий, удовлетворяющих обоим условиям.
    """
    # Условие 1: Браузер из списка целевых
    browser_condition: np.ndarray = np.isin(browsers, target_browsers)

    # Условие 2: Длительность сессии больше минимальной
    time_condition: np.ndarray = session_times > min_time_seconds

    # Комбинируем условия с помощью логического И (&)
    combined_condition: np.ndarray = browser_condition & time_condition

    # Получаем индексы с помощью np.where
    indices_tuple: tuple[np.ndarray, ...] = np.where(combined_condition)
    return indices_tuple[0]

# Пример использования (допустим, данные веб-аналитики)
browsers_data = np.array(['Chrome', 'Firefox', 'Safari', 'Chrome', 'Edge', 'Firefox', 'Chrome'])
times_data = np.array([45, 120, 30, 90, 50, 150, 75])
targets = ['Chrome', 'Firefox']
min_duration = 60

filtered_indices = find_indices_complex_condition(browsers_data, times_data, targets, min_duration)
# print(f"Индексы сессий из {targets} длительностью > {min_duration}s: {filtered_indices}")
# Вывод: Индексы сессий из ['Chrome', 'Firefox'] длительностью > 60s: [1 3 5 6]

Оптимизация производительности комбинированных методов

Использование комбинированных методов на основе np.where и np.isin полностью основано на векторизованных операциях NumPy. Это означает, что все вычисления выполняются над массивами целиком с использованием низкоуровневых оптимизаций, часто реализованных на C. Такой подход значительно быстрее, чем итерация по элементам массива в стандартных циклах Python, особенно для больших наборов данных. При работе с комбинированными условиями NumPy эффективно обрабатывает булевы маски, сводя несколько условий к одной итоговой булевой маске перед поиском индексов.

Альтернативные подходы и оптимизация

Использование циклов и условных операторов (для простых случаев)

Теоретически, можно найти индексы, итерируясь по массиву с помощью цикла for и проверяя условия для каждого элемента. Внутри цикла можно собирать индексы в список, если элемент соответствует критериям.

import numpy as np

def find_indices_with_loop(data_array: np.ndarray, target_value: int) -> list[int]:
    """
    Находит индексы элементов, равных целевому значению, используя цикл.
    (Пример для демонстрации, не рекомендуется для больших массивов).

    Args:
        data_array: Входной NumPy массив.
        target_value: Целевое значение для поиска.

    Returns:
        Список индексов элементов, равных целевому значению.
    """
    indices: list[int] = []
    for i in range(len(data_array)):
        if data_array[i] == target_value:
            indices.append(i)
    return indices

# Пример использования
scores = np.array([85, 92, 78, 85, 95, 88])
passing_score = 85

passing_indices_loop = find_indices_with_loop(scores.tolist(), passing_score) # Конвертируем в list для демонстрации
# print(f"Индексы элементов, равных {passing_score} (цикл): {passing_indices_loop}")
# Вывод: Индексы элементов, равных 85 (цикл): [0, 3]

Сравнение производительности различных методов

Использование циклов Python для обработки больших массивов NumPy крайне неэффективно. Каждая операция внутри цикла требует извлечения элемента из массива в стандартный тип данных Python, выполнения операции, а затем, возможно, сохранения результата обратно. Это приводит к значительному замедлению из-за накладных расходов на интерпретацию и управление памятью.

Векторизованные методы (np.where, np.isin, булева индексация) выполняют операции над всем массивом или его частями в скомпилированном коде (часто на C или Fortran), который лежит в основе NumPy. Это минимизирует накладные расходы и позволяет использовать преимущества оптимизаций процессора и многопоточности (если они поддерживаются NumPy для конкретной операции).

Таким образом, для любых нетривиальных размеров массивов векторизованные подходы будут на порядки быстрее циклов Python.

Рекомендации по выбору оптимального подхода в зависимости от размера массива и сложности условий

  • Для простых условий сравнения (больше, меньше, равно) или нескольких условий сравнения, связанных операторами И/ИЛИ: используйте np.where с булевой маской, полученной напрямую из сравнений (array > value, (cond1) & (cond2)). Это наиболее прямой и производительный способ.
  • Для поиска индексов элементов, которые принадлежат к определенному набору дискретных значений: используйте np.isin для создания булевой маски, а затем передайте ее в np.where или используйте булеву индексацию.
  • Для сложных условий, включающих как сравнения, так и проверку членства в наборе: комбинируйте np.isin и другие булевы маски, полученные сравнениями, используя логические операторы (&, |), и передавайте итоговую комбинированную маску в np.where.
  • Использование циклов Python для этой задачи следует избегать практически всегда при работе с массивами NumPy из-за их низкой производительности на больших данных. Они могут быть приемлемы только для очень маленьких массивов или в учебных целях.

Добавить комментарий