tkowt tkowt - 1 month ago 8
Python Question

How slicing inequality interval by numpy or scipy

I have numpy 2d array

A = [[1,1,1,1,1],[1,2,3,4,5]]


I want a function(A,B,axis) to assign interval,which results

([[1,1],[1,2]],[[1,1,1],[2,3,4]],[[1,1],[4,5]])


, effectively with interval position
B=[0,1,3,4]
and
axis=0
.In additional,it'll be better that lazy assignment each slice matrix like using generator because these matrixes can be very big size.
I know it easily accomplishment by for loop and yield iterator but I don't use loop for performance as possible.

Do you know the best way?

In my thought way,

def assign_interval(A,B,axis):
if axis == 0:
for i in range(len(B)-1):
yield A[:,B[i]:B[i+1]]
else:
for i in range(len(B)-1):
yield A[B[i]:B[i+1],:]


Edit:

I'm apologize my code didn't work.I was really busy today so that I couldn't inspect above code well and it's with the intention of dummy code for comprehension to process that I want.But,this code make terribly mistakes too much. Now, the code is revised and it'll be fine to work.

list(assign_interval(A,B,0))


results

[array([[1],
[1]]), array([[1, 1],
[2, 3]]), array([[1],
[4]])]


in my environment.

Answer

You could simply use np.split -

def assign_interval_split(A,B,axis):
    if axis == 0:
        return np.split(A[:,B[0]:B[-1]],B-B[0],axis=1)[1:-1]
    else:
        return np.split(A[B[0]:B[-1]],B-B[0],axis=0)[1:-1]