如何在不规则形状的numpy数组中选取各子数组首尾元素

给定如下numpy数组：

import numpy as np

v = np.array([
    np.array([1, 1]),
    np.array([1, 2]),
    np.array([1, 3]),
    np.array([1, 4]),
    np.array([1, 5]),
    np.array([2, 1]),
    np.array([2, 2]),
    np.array([2, 3]),
    np.array([3, 1]),
    np.array([3, 2]),
    np.array([3, 3]),
    np.array([3, 4]),
    np.array([4, 1]),
    np.array([4, 2]),
    np.array([4, 3]),
    np.array([4, 4]),
    np.array([4, 5]),
    np.array([4, 6]),
])

原本通过np.split(v[:, 1], np.unique(v[:, 0], return_index=True)[1][1:])得到分组后的数组列表，现在需要直接获取每个分组数组的首尾元素，要求仅用原生numpy实现，不使用列表推导、函数或循环。

可以利用np.unique的return_index和return_counts参数，直接定位每个分组的首尾元素索引，无需先拆分数组：

import numpy as np

# 获取分组键、起始索引、每组元素个数
unique_keys, start_indices, group_counts = np.unique(v[:, 0], return_index=True, return_counts=True)

# 计算每组最后一个元素的索引
end_indices = start_indices + group_counts - 1

# 提取首尾元素并组合成结果
result = np.vstack([v[start_indices, 1], v[end_indices, 1]]).T

代码解释 #

• np.unique(v[:,0], return_index=True, return_counts=True)：返回分组的唯一键、每个分组第一个元素的索引、每个分组的元素数量。
• end_indices = start_indices + group_counts - 1：通过起始索引加元素个数减1，得到每个分组最后一个元素的索引。
• np.vstack([...]).T：将首尾元素分别按行堆叠后转置，得到每个分组对应首尾元素的二维数组（每行对应一个分组的首尾）。

输出结果 #

运行上述代码后，result的值为：

array([[1, 5],
       [1, 3],
       [1, 4],
       [1, 6]])

如果需要保持列表形式（每个元素是包含首尾的数组），可以用原生numpy操作实现：

result_list = np.split(result, len(result))
# 输出：[array([1,5]), array([1,3]), array([1,4]), array([1,6])]

如果分组后的数组长度不固定，但仍想避免显式循环，也可以用np.take配合索引数组直接提取：

# 合并首尾索引
indices = np.concatenate([start_indices, end_indices])
# 提取对应元素并重新分组
elements = v[indices, 1].reshape(-1, 2)

效果和之前的解决方案一致，同样是纯numpy操作。

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