统计三维整数张量每行的唯一k维向量数量及获取唯一向量

问题背景 #

我们有一个形状为(m, n, k)的三维整数张量X，可以把它理解成一个(m, n)的矩阵——矩阵里的每个元素都是一个k维整数向量。需要实现两个核心功能：

1. 统计每行中这类唯一向量的数量
2. 提取每行对应的所有唯一向量

举个实际例子：

import numpy as np

X = np.array([[[0, 1, 2], [0, 1, 2], [1, 2, 3], [1, 2, 3]],
              [[3, 4, 5], [4, 5, 6], [5, 6, 7], [6, 7, 8]]])
print(X.shape)  # 输出 (2, 4, 3)

我们需要得到：

• 统计结果：[2, 4]（第一行有2个唯一向量，第二行有4个）
• 提取的唯一向量集合：

[[[0, 1, 2], [1, 2, 3]],
 [[3, 4, 5], [4, 5, 6], [5, 6, 7], [6, 7, 8]]]

解决方案 #

1. 统计每行唯一向量的数量

你提到的部分向量化方案已经非常实用，这里拆解下逻辑：

import numpy as np

count_unique = lambda X: [len(np.unique(row, axis=0)) for row in X]

• 先遍历张量的每一行row（每一行是形状为(n, k)的二维数组，对应n个k维向量）
• np.unique(row, axis=0)会对这一行的向量按行维度去重，得到该行所有的唯一向量
• 取去重后数组的长度，就是该行唯一向量的数量

运行示例验证：

print(count_unique(X))  # 输出 [2, 4]

2. 提取每行的唯一向量

同样用列表推导式结合np.unique实现，直接返回每行的唯一向量集合：

get_unique = lambda X: [np.unique(row, axis=0) for row in X]

• 逻辑和统计数量类似，只是不需要取长度，直接保留np.unique返回的去重后向量数组
• 返回结果是一个列表，每个元素对应原张量一行的唯一向量数组

运行示例验证：

result = get_unique(X)
print(result[0])  # 输出 [[0 1 2], [1 2 3]]
print(result[1])  # 输出 [[3 4 5], [4 5 6], [5 6 7], [6 7 8]]

补充说明 #

如果你的张量规模特别大，逐行处理的方式虽然直观，但可以考虑用更高效的全向量化操作（比如利用np.unique的return_inverse参数结合分组统计），不过对于大多数日常场景，上面的实现已经足够清晰且高效了。

内容的提问来源于stack exchange，提问作者H.Rappeport