且构网

分享程序员开发的那些事...
且构网 - 分享程序员编程开发的那些事

如何在NumPy中将N-D切片转换为索引?

更新时间:2023-01-22 20:34:13

感谢 np.mgrid .

在线试用!

import numpy as np

def nd_slice_to_indexes(nd_slice):
    grid = np.mgrid[{tuple: nd_slice, slice: (nd_slice,)}[type(nd_slice)]]
    return tuple(grid[i].ravel() for i in range(grid.shape[0]))
    
print(nd_slice_to_indexes(np.s_[1 : 3]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2, 8 : 14 : 3]))
# (array([1, 2]),)
# (array([1, 1, 1, 2, 2, 2]), array([5, 7, 9, 5, 7, 9]))
# (array([1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2]), array([5, 5, 7, 7, 9, 9, 5, 5, 7, 7, 9, 9]), array([ 8, 11,  8, 11,  8, 11,  8, 11,  8, 11,  8, 11]))