如何查找3Dnumpy数组每列中连续1的组数

import numpy as np

def count_consecutive_ones(arr):
    # 初始化计数器
    counts = []
    
    # 遍历每一列
    for col in range(arr.shape[2]):
        current_count = 0
        max_count = 0
        
        # 遍历当前列的每个元素
        for row in range(arr.shape[0]):
            if arr[row, :, col].sum() == arr.shape[1]:  # 检查当前行在该列是否全为1
                current_count += 1
                max_count = max(max_count, current_count)
            else:
                current_count = 0
        
        counts.append(max_count)
    
    return counts

# 示例3D数组
example_array = np.array([
    [[1, 1, 0], [1, 1, 0]],
    [[1, 1, 1], [1, 1, 1]],
    [[0, 0, 0], [0, 0, 0]],
    [[1, 1, 1], [1, 1, 1]]
])

# 计算每列中连续1的组数
result = count_consecutive_ones(example_array)
print("每列中连续1的组数:", result)