N-мерные индексы уникальных значений в ndarray

Question

У меня есть двумерный массив Numpy, содержащий значения от 0 до n.Я хочу получить список длины n, такой, что i-й элемент этого списка является массивом всех индексов со значением i + 1 (0 исключено).

Например, для ввода

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

Я ожидаю получить

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

Я нашел этот связанный вопрос: Получить список всех индексов повторяющихся элементов в массиве numpy , которыйможет быть полезным, но я надеялся найти более прямое решение, которое не требует выравнивания и сортировки массива, и которое будет максимально эффективным.

Answer 1

Вот векторизованный подход, который работает для массивов произвольного количества измерений.Идея этого решения заключается в расширении функциональности метода return_index в np.unique и возвращении массива массивов, каждый из которых содержит одномерные индексы, где появляются уникальные значения в данном массиве numpy.

Для более компактного решения я определил следующую функцию вместе с некоторыми пояснениями на разных этапах.

---

Общее решение

def ndix_unique(x):
    """
    Finds the indices where each unique
    value in x is found
    ----------
    x: np.array
       Can have multiple dimensions
    Returns
    -------
    - 1darray of sorted unique values
    - Array of arrays. Each array contains the indices where a
      given value in x is found, where the values are sorted
    """
    # flattens x. Both will be necessary
    x_ = x.ravel()
    # Sort y using np.argsort
    ix_ = np.argsort(x_)
    # Finds the indices in x_ that result in the
    # unique array. Used later on to split
    u, s = np.unique(x_[ix_], return_index=True)
    # Mapping of the 1D indices to n-dimensional
    # indices taking the shape of x as a reference
    ix_ndim = np.unravel_index(ix_, x.shape)
    # Stack these as columns 
    ix = np.column_stack(ix_ndim) if x.ndim > 1 else ix_
    # Split the nD coordinates using the indices in s
    # i.e. where the changes of values take place
    return u, np.split(ix, s[1:])

---

Примеры

Давайте начнем с проверки результатов с помощью предложенного ndarray:

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

vals, ixs = ndix_unique(a)

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

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

---

Давайте попробуем в другом случае:

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

vals, ixs = ndix_unique(a)

print(vals)
array([1, 2, 3, 4])

print(ixs)
array([array([[0, 0],
              [0, 1],
              [1, 2],
              [2, 2]]),
       array([[1, 0],
              [1, 1]]), 
       array([[2, 0],
              [2, 1]]),
       array([[0, 2]])], dtype=object)

---

Теперь попробуем с массивом 1D:

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

vals, ixs = ndix_unique(a)

print(vals)
array([1, 3, 4, 5])

print(ixs)
array([array([0]), array([3, 4]), array([2]), array([1])], dtype=object)

---

И, наконец, еще один пример с 3D ndarray:

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

vals, ixs = ndix_unique(a)

print(vals)
array([1, 2, 3, 4])

print(ixs)
array([array([[0, 0, 0],
              [0, 0, 1]]),
       array([[0, 0, 2],
              [1, 0, 0]]), 
       array([[1, 0, 1]]),
       array([[1, 0, 2]])], dtype=object)

Answer 2

Вы можете сначала получить ненулевые элементы в вашем массиве, а затем использовать argwhere в понимании списка, чтобы получить отдельный массив для каждого ненулевого элемента.Здесь np.unique(arr[arr!=0]) даст вам ненулевые элементы, по которым вы можете выполнить итерацию для получения индексов.

arr = np.array([[1, 0, 1],
            [2, 2, 0]])

indices = [np.argwhere(arr==i) for i in np.unique(arr[arr!=0])]
# [array([[0, 0],
#         [0, 2]]), array([[1, 0],
#         [1, 1]])]