如何为numpy数组添加列，行值变化时递增计数且起始值为3

实现思路 #

遍历数组的所有行，记录上一行的内容与当前计数值，初始计数值为3。每遇到与上一行内容不同的行时，计数值加1，最终将计数值作为新列追加到当前行末尾即可。

纯Python实现（无需第三方依赖） #

# 原始数组
raw_arr = [
    [1, 2, 1, 0, 2, 0],
    [1, 2, 1, 0, 2, 0],
    [1, 2, 1, 0, 2, 0],
    [1, 2, 1, 0, 2, 0],
    [0, 1, 2, 1, 0, 0],
    [0, 1, 2, 1, 0, 0],
    [0, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0],
    [0, 0, 0, 0, 1, 0],
    [0, 0, 0, 0, 0, 1]
]

result = []
if raw_arr:
    count = 3
    last_row = raw_arr[0]
    # 处理第一行
    result.append(last_row + [count])
    # 遍历剩余行
    for row in raw_arr[1:]:
        if row != last_row:
            count += 1
            last_row = row
        result.append(row + [count])

# 输出结果
for line in result:
    print(line)

Numpy实现（适合处理大规模数组） #

import numpy as np

raw_arr = np.array([
    [1, 2, 1, 0, 2, 0],
    [1, 2, 1, 0, 2, 0],
    [1, 2, 1, 0, 2, 0],
    [1, 2, 1, 0, 2, 0],
    [0, 1, 2, 1, 0, 0],
    [0, 1, 2, 1, 0, 0],
    [0, 0, 1, 0, 1, 0],
    [0, 0, 0, 1, 1, 0],
    [0, 0, 0, 0, 1, 0],
    [0, 0, 0, 0, 0, 1]
])

# 计算相邻行差异
row_diff = np.any(raw_arr[1:] != raw_arr[:-1], axis=1)
# 生成计数列
count_col = np.concatenate(([3], 3 + np.cumsum(row_diff)))
# 拼接列得到最终结果
result = np.column_stack((raw_arr, count_col))

print(result)

两种实现运行后都能得到你期望的输出结果。

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