Início / coding / Perguntas / 79227441
Accepted
Falcondance
Asked: 2024-11-27 00:02:39 +0800 CST

Argsort mascarado eficiente em Numpy

  • 772

Eu tenho um array numpy como este:

arr = np.array([
    [1, 2, 3],
    [4, -5, 6],
    [-1, -1, -1]
])

E eu gostaria de argsortá-lo, mas com uma arr <= 0máscara. A saída deve ser:

array([[0, 1, 2],
       [0, 2],       # (Note that the indices are still relative to original un-masked array)
       []])

Entretanto, a saída que obtenho usando np.ma.argsort()é:

array([[0, 1, 2],
       [0, 2, 1],
       [0, 1, 2]])

A abordagem precisa ser muito eficiente porque o array real tem milhões de colunas. Estou pensando que isso precisa ser uma síntese de algumas operações, mas não sei quais.

python

1 respostas

  1. Best Answer

    A np.whereabordagem:

    Matriz de entrada

    arr = np.array([
    [1, 2, 3],
    [4, -5, 6],
    [-1, -1, -1]
])

    Máscara de elementos válidos

    mask = arr > 0

    Pré-alocar o resultado como uma matriz de objetos para conter índices de comprimento variável

    result = np.empty(arr.shape[0], dtype=object)

    Máscara eficiente argsortpara cada linha

    for i in range(arr.shape[0]):
    valid_indices = np.where(mask[i])[0]  # Get indices of valid (masked) elements
    result[i] = valid_indices[np.argsort(arr[i, valid_indices])]  # Sort valid indices by their values

    Saída:

    [array([0, 1, 2]) array([0, 2]) array([], dtype=int64)]

    A np.flatnonzeroabordagem:

    Uma abordagem mais otimizada usando operações vetorizadas:

    def optimized_masked_argsort(arr, mask):
    result = np.empty(arr.shape[0], dtype=object)
    for i in range(arr.shape[0]):
        row = arr[i]
        valid_indices = np.flatnonzero(mask[i])  # Faster than np.where(mask[i])[0]
        valid_values = row[valid_indices]
        sorted_order = np.argsort(valid_values)
        result[i] = valid_indices[sorted_order]
    return result

    Comparação:

    Timings for given example:
np.where Time: 0.000034 seconds
np.flatnonzero Time: 0.000017 seconds

Timings for larger array (1000 rows):
np.where Time: 0.001856 seconds
np.flatnonzero Time: 0.001754 seconds

    Tentei alguns outros métodos, mas eles não foram eficientes.

    • 1

relate perguntas