Булев индекс для многомерного массива

Question

Как мне создать логический индекс для многомерного массива? Я также пробовал данные [метки == 0,:,:]

Данные имеют форму (6000, 28, 28)

Этикетки имеют форму (6000, 1)

new = data[labels == 0]

IndexError: логический индекс не соответствует индексированному массиву по измерению 1; размерность 28, но соответствующее логическое измерение 1

Я могу запустить следующее

new= data[np.where(labels == c)]

но моя возвращенная форма имеет (num_matching_criteria, 28)

Answer 1

Я смог сделать следующее

sel = np.array(labels == c)[:,0]
slice = data[sel]

Answer 2

Вы используете массив numpy? Это работает для меня, когда и данные, и метки являются массивами.

>>> import numpy
>>> a = numpy.arange(24).reshape((4,3,2))
>>> a
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],
       [[ 6,  7],
        [ 8,  9],
        [10, 11]],
       [[12, 13],
        [14, 15],
        [16, 17]],
       [[18, 19],
        [20, 21],
        [22, 23]]])
>>> labels = numpy.array([1,0,1, 0])
>>> a[labels > 0]
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],
       [[12, 13],
        [14, 15],
        [16, 17]]])
>>> a[labels == 0]
array([[[ 6,  7],
        [ 8,  9],
        [10, 11]],
       [[18, 19],
        [20, 21],
        [22, 23]]])