Получить индексы для разделения NumPy массив

Question

Допустим, у меня есть массив NumPy:

x = np.array([3, 9, 2, 1, 5, 4, 7, 7, 8, 6])

Если я суммирую этот массив, я получу 52. Что мне нужно, так это способ разбить его, начиная слева направо, примерно на 10000 * кусков, где пользователь выбирает n. По сути, расколы происходят жадным способом. Таким образом, для некоторого числа чанков n первые n - 1 чанков должны каждый составлять не менее 52/n, и они должны быть последовательными индексами, взятыми слева направо.

Итак, если n = 2 тогда первый блок будет состоять из первых 7 элементов:

chunk[0] = x[:7]  # [3, 9, 2, 1, 5, 4, 7], sum = 31
chunk[1] = x[7:]  # [7, 8, 6], sum = 21

Обратите внимание, что первый блок не будет состоять только из первых 6 элементов, поскольку сумма будет 24, что меньше 52/2 = 26. Кроме того, обратите внимание, что количество элементов в каждом чанке может варьироваться при условии соблюдения критерия суммы. Наконец, вполне нормально, чтобы последний блок не был близок к 52/2 = 26, так как другой блок (ы) может занять больше.

Однако мне нужен выходной массив из двух столбцов, который содержит начальный индекс в первом столбце и (исключительный) конечный индекс во втором столбце:

[[0, 7],
 [7, 10]]

Если n = 4, то первые 3 блока должны составлять не менее 52/4 = 13. выглядят так:

chunk[0] = x[:3]  # [3, 9, 2], sum = 14
chunk[1] = x[3:7]  # [1, 5, 4], sum = 17
chunk[2] = x[7:9]  # [7, 8], sum = 15
chunk[3] = x[9:]  # [6], sum = 6

И вывод, который мне нужен, будет:

[[0, 3],
 [3, 7],
 [7, 9],
 [9, 10]

Итак, один наивный подход, использующий циклы for, может быть:

ranges = np.zeros((n_chunks, 2), np.int64)
ranges_idx = 0
range_start_idx = start

sum = 0
for i in range(x.shape[0]):
    sum += x[i]
    if sum > x.sum() / n_chunks:
        ranges[ranges_idx, 0] = range_start_idx
        ranges[ranges_idx, 1] = min(
                i + 1, x.shape[0]
            )  # Exclusive stop index
        # Reset and Update
        range_start_idx = i + 1
        ranges_idx += 1
        sum = 0
# Handle final range outside of for loop
ranges[ranges_idx, 0] = range_start_idx
ranges[ranges_idx, 1] = x.shape[0]
if ranges_idx < n_chunks - 1:
    left[ranges_idx:] = x.shape[0]

return ranges

Я ищу более хорошее векторизованное решение.

Answer 1

Я нашел вдохновение в подобном вопросе, на который был дан ответ :

def func(x, n):
    out = np.zeros((n, 2), np.int64)
    cum_arr = x.cumsum() / x.sum()
    idx = 1 + np.searchsorted(cum_arr, np.linspace(0, 1, n, endpoint=False)[1:])
    out[1:, 0] = idx  # Fill the first column with start indices
    out[:-1, 1] = idx  # Fill the second column with exclusive stop indices
    out[-1, 1] = x.shape[0]  # Handle the stop index for the final chunk
    return out

Обновление

Чтобы охватить патологический случай, нам нужно чтобы быть немного более точным и сделать что-то вроде:

def func(x, n, truncate=False):
    out = np.zeros((n_chunks, 2), np.int64)
    cum_arr = x.cumsum() / x.sum()
    idx = 1 + np.searchsorted(cum_arr, np.linspace(0, 1, n, endpoint=False)[1:])
    out[1:, 0] = idx  # Fill the first column with start indices
    out[:-1, 1] = idx  # Fill the second column with exclusive stop indices
    out[-1, 1] = x.shape[0]  # Handle the stop index for the final chunk

    # Handle pathological case
    diff_idx = np.diff(idx)
    if np.any(diff_idx == 0):
        row_truncation_idx = np.argmin(diff_idx) + 2
        out[row_truncation_idx:, 0] = x.shape[0]
        out[row_truncation_idx-1:, 1] = x.shape[0]
        if truncate:
            out = out[:row_truncation_idx]

    return out

Answer 2

Вот решение, которое не перебирает все элементы:

def fun2(array, n):
    min_sum = np.sum(array) / n
    cumsum = np.cumsum(array)
    i = -1
    count = min_sum
    out = []
    while i < len(array)-1:
        j = np.searchsorted(cumsum, count) 
        out.append([i+1, j+1])
        i = j 
        if i < len(array):
            count = cumsum[i] + min_sum
    out[-1][1] -= 1
    return np.array(out)

Для двух тестовых случаев оно дает ожидаемые вами результаты. НТН