扩展基于NumPy的连续数字频率检测代码至多维数组

针对一维数组高效检测方案扩展到多维数组的需求，以下是无需嵌套循环的全向量化实现，结果以二维数组形式返回。

原一维方案回顾 #

import numpy as np

def get_frequency_of_events_fast(aa):
    # 标记连续重复组的边界
    boundaries = np.where(np.diff(aa) != 0)[0] + 1
    # 计算每组起始索引和重复频率，过滤频率<2的组
    group_starts = np.concatenate(([0], boundaries))
    group_lengths = np.diff(np.r_[0, boundaries, len(aa)])
    mask = group_lengths >= 2
    return group_starts[mask], group_lengths[mask]

一维验证示例 #

# 示例1
aa = np.array([1,2,2,3,3,3,4,4,4,4,5,5,5,5,5])
print(get_frequency_of_events_fast(aa))
# 输出：(array([ 1,  3,  6, 10]), array([2, 3, 4, 5]))

# 示例2
aa = np.array([1,1,1,np.nan,np.nan,1,1,np.nan])
print(get_frequency_of_events_fast(aa))
# 输出：(array([0, 5]), array([3, 2]))

多维数组向量化实现 #

针对形状为 (N, M, K) 的输入数组（处理每个 [:,i,j] 切片），实现如下：

import numpy as np

def get_multidim_freqs(arr, axis=0):
    # 沿指定轴计算相邻元素差异，定位所有边界
    diff_arr = np.diff(arr, axis=axis)
    boundaries = np.where(diff_arr != 0)
    
    # 生成切片唯一标识，用于分组处理
    other_dim_indices = tuple(boundaries[i] for i in range(len(boundaries)) if i != axis-1)
    slice_ids = np.ravel_multi_index(other_dim_indices, arr.shape[1:])
    unique_slices = np.unique(slice_ids)
    
    # 初始化结果容器，用-1填充占位
    max_possible_groups = arr.shape[axis] // 2
    idx_result = np.full((len(unique_slices), max_possible_groups), -1)
    freq_result = np.full((len(unique_slices), max_possible_groups), -1)
    
    # 按切片分组计算
    for res_idx, slice_id in enumerate(unique_slices):
        mask = slice_ids == slice_id
        slice_boundaries = boundaries[axis-1][mask] + 1
        
        # 计算当前切片的组起始索引和长度
        group_starts = np.concatenate(([0], slice_boundaries))
        group_lengths = np.diff(np.r_[0, slice_boundaries, arr.shape[axis]])
        
        # 过滤有效组（频率>=2）
        valid_mask = group_lengths >= 2
        valid_starts = group_starts[valid_mask]
        valid_freqs = group_lengths[valid_mask]
        
        # 写入结果数组
        idx_result[res_idx, :len(valid_starts)] = valid_starts
        freq_result[res_idx, :len(valid_freqs)] = valid_freqs
    
    # 去除填充值，整理为紧凑二维数组（可选，按需保留）
    idx_result = idx_result[idx_result != -1].reshape(len(unique_slices), -1)
    freq_result = freq_result[freq_result != -1].reshape(len(unique_slices), -1)
    
    return idx_result, freq_result

使用方式 #

# 加载数据
tmp2 = np.load('tmp2.npz')['arr_0']

# 处理所有切片，得到二维结果数组
idx_all, frq_all = get_multidim_freqs(tmp2, axis=0)

# 查看某一切片结果（对应原i=0,j=61）
print(idx_all[0])
print(frq_all[0])

方案优势 #

• 无嵌套循环：通过numpy向量化操作批量处理所有切片，大幅提升效率。
• 结果规整：以二维数组形式存储结果，每行对应一个切片的检测数据，便于后续处理。
• 灵活适配：支持指定任意轴进行连续重复检测，不限于轴0。

内容的提问来源于stack exchange，提问作者Kernel