根据输入偏移值拆分数组，但在同一块中保留重复

Question

给定一个索引列表（偏移值），根据该列表拆分一个 numpy 数组，我想对其进行调整，以便在重复值上不会发生拆分。 这意味着重复值将仅在一个块中。

我已经编写了以下代码，它给出了结果，但我并不为此感到非常自豪。我想留在 numpy 世界中，并尽可能多地使用向量化的 numpy 函数。 但要检查索引（偏移值），我使用for 循环，并将结果存储在列表中。

你知道如何矢量化第二部分吗？

如果这有帮助，ar 是一个有序数组。 （我没有在下面的代码中使用此信息）。

import numpy as np
import vaex as vx

ar = np.array([8,8,8,10,11,11,13,14,15,15,18,19,20,21,22,22,22,22,22,22])
offsets = np.array([0,2,4,9,11,13,15,len(ar)])

_, unique_ind = np.unique(ar, return_index=True, return_inverse=False)
dup_ind = np.diff(unique_ind, append=len(ar))
dup_start = unique_ind[dup_ind > 1]
dup_end = dup_start + dup_ind[dup_ind > 1]
print(f'initial offsets: {offsets}')
#print(f'dup start: {dup_start}')
#print(f'dup end: {dup_end}')

temp = []
for off in offsets:
    for ind in range(len(dup_start)):
        if off > dup_start[ind] and off < dup_end[ind]:
            off = dup_start[ind]
            break
    temp.append(off)

# Remove duplicates
offsets = list(dict.fromkeys(temp))
print(f'modified offsets: {offsets}')

结果

initial offsets: [ 0  2  4  9 11 13 15 20]
modified offsets: [0, 4, 8, 11, 13, 14, 20]

Answer 1

根据@richardec 的回答，这是一个可能的解决方案，以防它可以帮助其他人。

import numpy as np

def adjust_offsets(values, offsets) -> np.ndarray:
    # Creates bins describing when duplicates start and when they end.
    _, indexes, counts = np.unique(values, return_index=True, return_counts=True)
    duplicates = counts>1
    dup_starts = indexes[duplicates]
    dup_ends = dup_starts + counts[duplicates]
    # Zipping starts and ends.
    # Bins of duplicates are interlayed with bins of unique values.
    bins = np.dstack((dup_starts,dup_ends)).flatten()
    # Set start of bin for offsets falling in a bin.
    off_binned = np.digitize(offsets, bins, right=True)
    return np.unique(np.where((off_binned%2)==0,offsets,bins[off_binned-1]))

# Test setup 1 / starting with no duplicates
# idx:         0,1,2,3,4, 5, 6, 7, 8, 9,10,11,12,13,14,15,16,17,18,19,20,21
ar = np.array([4,5,8,8,8,10,11,11,13,14,15,15,18,19,20,21,22,22,22,22,22,22])
offsets = np.array([0,2,4,9,11,14,18,len(ar)-1])
off_adjusted = adjust_offsets(ar, offsets)
ref = np.array([ 0,  2,  9, 10, 14, 16])
assert np.array_equal(off_adjusted, ref)

# Test setup 2 / starting with duplicates
# idx:         0,1,2, 3, 4, 5, 6, 7, 8, 9,10,11,12,13,14,15,16,17,18,19
ar = np.array([8,8,8,10,11,11,13,14,15,15,18,19,20,21,22,22,22,22,22,22])
offsets = np.array([0,2,4,9,11,13,15,len(ar)-1])
off_adjusted = adjust_offsets(ar, offsets)
ref = np.array([ 0,  4,  8, 11, 13, 14])
assert np.array_equal(off_adjusted, ref)

我使用了 2 个测试，因为如果数组不是以重复开头，我想出的以前的版本将无法工作。

2句话，算法遵循以下逻辑：

• 如果偏移值（索引）落在由 [start of duplicate, end of duplicate] 定义的 bin 中，则将其重置为“start of duplicate”。
• 如果不是，则未修改

可能更快，正如我从这个相关的SO answer 中看到的那样，np.diff 比np.unique 更快，我认为可以使用np.diff 获得相同类型的信息。我开始评估这种可能性，但它对我来说太复杂了。

Answer 2

您可以使用np.digitize 将偏移量限制在 bin 中：

ar = np.array([8,8,8,10,11,11,13,14,15,15,18,22,22,22,22,22,22])
offsets = np.array([0,2,4,9,13,15,len(ar)])

_, unique_ind = np.unique(ar, return_index=True, return_inverse=False)
dup_ind = np.diff(unique_ind, append=len(ar))
dup = np.append(unique_ind[dup_ind > 1], len(ar))

offsets = dup[np.digitize(offsets, dup) - 1]

输出：

>>> offsets
array([ 0,  0,  4,  8, 11, 11, 17])

>>> np.unique(offsets)
array([ 0,  4,  8, 11, 17])

它很可能被压缩成更少的代码（可能是 1 或 2 行），但我认为你真的只是希望消除循环，所以我想提交我发现的内容。