Поиск случайного индекса с конкретным значением в большом массиве

Question

У меня очень большой двумерный массив NumPy (значения ~ 5e8).Я пометил этот массив с помощью scipy.ndimage.label Затем я хочу найти случайный индекс уплощенного массива, который содержит каждую метку.Я могу сделать это с:

import numpy as np
from scipy.ndimage import label

base_array = np.random.randint(0, 5, (100000, 5000))
labeled_array, nlabels = label(base_array)
for label_num in xrange(1, nlabels+1):
    indices = np.where(labeled_array.flat == label_num)[0]
    index = np.random.choice(indices)

Но, с таким большим массивом, это медленно.Я также попытался заменить np.where на:

indices = np.argwhere(labeled_array.flat == label).squeeze()

и обнаружил, что он медленнее.У меня есть подозрение, что логическое маскирование является медленной частью.Есть ли способ ускорить это или лучший способ сделать это?Я скажу, что в моем реальном приложении массив довольно редок с заполнением около 25%, хотя у меня нет опыта работы с редкими функциями массива scipy.

Answer 1

Ваше подозрение, что маскировка отдельно для каждой метки является дорогой, является правильным, потому что независимо от того, как вы это сделаете, маскировка всегда будет O (n).

Мы можем обойти это путем сортировки по меткам, а затем случайным образомвыбор из каждого блока одинаковых меток.

Поскольку метки представляют собой целочисленный диапазон, мы можем получить аргумент дешевле, чем np.argsort, используя некоторые разреженные матричные механизмы, доступные в scipy.

В качестве моей машиныУ меня не так много оперативной памяти. Мне пришлось немного сократить ваш пример (фактор 4).Затем он запускается примерно через 5 секунд.

import numpy as np
from scipy.ndimage import label
from scipy import sparse

def multi_randint(bins):
    """draw one random int from each range(bins[i], bins[i+1])"""
    high = np.diff(bins)
    n = high.size
    pick = np.random.randint(0, 1<<30, (n,))
    reject = np.flatnonzero(pick + (1<<30) % high >= (1<<30))
    while reject.size:
        npick = np.random.randint(0, 1<<30, (reject.size,))
        rejrej = npick + (1<<30) % sizes[reject] >= (1<<30)
        pick[reject] = npick
        reject = reject[rejrej]
    return bins[:-1] + pick % high

# build mock data, note that I had to shrink by 4x b/c memory
base_array = np.random.randint(0, 5, (50000, 2500), dtype=np.int8)
labeled_array, nlabels = label(base_array)

# build auxiliary sparse matrix
h = sparse.csr_matrix(
    (np.ones(labeled_array.size, bool), labeled_array.ravel(),
     np.arange(labeled_array.size+1, dtype=np.int32)),
    (labeled_array.size, nlabels+1))
# conversion to csc argsorts the labels (but cheaper than argsort)
h = h.tocsc()
# draw
result = h.indices[multi_randint(h.indptr)]

# check result
assert len(set(labeled_array.ravel()[result])) == nlabels+1