python - 获取分割 NumPy 数组的索引

Question

假设我有一个 NumPy 数组:

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

如果我对这个数组求和，我会得到52。我需要的是一种将其从左到右分成大约 n 个 block 的方法，其中 n 由用户选择。本质上， split 是以贪婪的方式发生的。因此，对于一定数量的 block n，前 n - 1 block 的每个总和必须至少为 52/n 并且它们必须是从左到右取连续索引。

因此，如果n = 2，那么第一个 block 将包含前 7 个元素:

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

请注意，第一个 block 不会仅包含前 6 个元素，因为总和将为 24，小于 52/2 = 26。另请注意，只要满足总和标准，每个 block 中的元素数量就可以变化。最后，最后一个 block 不接近 52/2 = 26 是完全可以的，因为其他 block 可能需要更多。

但是，我需要的输出是一个两列数组，其中包含第一列中的开始索引和第二列中的(独占)停止索引:

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

如果n = 4，则前 3 个 block 的总和至少需要 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

我在 similar question that was answered 中找到了灵感:

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