Numpy: векторизация присвоения значений в зависимости от условий

Question

Рассмотрим следующую функцию:

import numpy
import scipy.stats


def return_category(values, categories):
    n = len(categories)

    result = numpy.empty(values.shape, dtype='U25')

    boundaries = scipy.stats.norm.ppf(numpy.arange(0, n+1, 1)/n)
    for i, category in enumerate(categories):
        a, b = boundaries[i], boundaries[i + 1]
        numpy.putmask(result, (values < b) & (values >= a), category)

    return result


print(return_category(numpy.array([0.1, -100, 100, 0.44]), ['a', 'b', 'c']))
# ['b' 'a' 'c' 'c']

т.е. он назначает категорию из списка категорий в зависимости от того, где находится значение, так что каждая категория одинаково вероятна, если values взяты из нормального распределения (0, 1).

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

Эту проблему можно сформулировать в более общем виде следующим образом: существует карта M={I1: c1, I2: c2, ...}, где Ii - это такой интервал, что объединение всех интервалов равно ]-inf,inf[, их пересечение пусто, а ci - это категория , Учитывая массив значений [a1, a2, ..., aM], создайте новый массив

[
 M[Ii such that a1 in Ii],
 M[Ii such that a2 in Ii], 
 ...
 M[Ii such that aM in Ii],
]

В приведенном выше конкретном случае интервалы составляют scipy.stats.norm.ppf(numpy.arange(0, n+1, 1)/n)

Answer 1

Я думаю, что это может сделать то, что вы хотите:

import numpy 
import scipy.stats


def return_category(values, categories):
    n = len(categories)
    categories = numpy.array(categories)
    result = numpy.empty(values.shape, dtype='U25')
    boundaries = scipy.stats.norm.ppf(numpy.arange(0, n+1, 1)/n)
    # array of "left" boundaries
    bndrs0 = boundaries[:-1]
    # array of "right" boundaries
    bndrs1 = boundaries[1:]
    # build an array such that the j-th column in the
    # i-th row is True if the j-th column of values is in the i-th category
    whereCat = numpy.where(numpy.logical_and(values>=numpy.tile(bndrs0, (values.size,1)).T, values < numpy.tile(bndrs1, (values.size,1)).T))
    # broadcast categories to the corresponding rows
    sortedCats = numpy.take_along_axis(categories, whereCat[0],0)
    # place categories in the correct column
    numpy.put_along_axis(result,whereCat[1],sortedCats,0)
    return result


print(return_category(numpy.array([0.1, -100, 100, 0.44]), ['a', 'b', 'c']))
# ['b' 'a' 'c' 'c']