如何获取NumPy二维数组元素沿指定轴的位置索引，供matplotlib bar3d使用

实现方案 #

你需要的是指定轴的索引广播匹配原数组形状，用NumPy内置的广播机制就能一行搞定，不需要循环：

1. 直接生成你要的axis=1方向的索引

import numpy as np
a = np.arange(1,7).reshape(2,3)
# 沿axis=1（列方向）取位置索引
res = np.broadcast_to(np.arange(a.shape[1]), a.shape)
print(res)

输出完全符合你的预期：

array([[0, 1, 2],
       [0, 1, 2]])

2. 通用的任意轴索引生成方法

如果需要灵活指定轴，可以封装成通用写法：

def get_axis_pos_index(arr, target_axis):
    arr_shape = arr.shape
    # 生成目标轴的索引序列
    idx_seq = np.arange(arr_shape[target_axis])
    # 调整索引序列的维度适配广播规则
    expand_dims = [np.newaxis] * len(arr_shape)
    expand_dims[target_axis] = slice(None)
    # 广播到和原数组同形状，无额外内存拷贝
    return np.broadcast_to(idx_seq[tuple(expand_dims)], arr_shape)

调用示例：

• 沿列方向（axis=1）取索引：get_axis_pos_index(a, 1) 输出和上面一致
• 沿行方向（axis=0）取索引：get_axis_pos_index(a, 0) 输出为 [[0 0 0],[1 1 1]]

这个方法性能很高，生成的是视图而非拷贝数组，完全满足bar3d方法的参数输入要求。

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