CODEWITHSUNDEEP
pythonarraysnumpy
konstantin
konstantin

Reputation: 893

Split a 2D numpy array where uneven splits are possible

I have an 2D numpy arrays in python which correspond to image that are calculated in a for-loop. The size of the arrays are Nx40. I want in each step of the loop to split the initial that arrays into rectangular arrays of size 40x40 (approximately). In case that N cannot divided with 40 then the last image should contain the remaining of the division. Therefore, for example of 87x40 should be (40x40 and 47x40). What I did so far:

div_num = spec.shape[0] / spec.shape[1]
remaining = spec.shape[0] % spec.shape[1]

lista = []
for i in range(1, div_num+1):
     img = spec[((i-1)*40):(i*40)][0:40]
     lista.append(img)

How can I add the remaining rows in the last image?

Upvotes: 3

Views: 1773

Answers (1)

cs95
cs95

Reputation: 402283

You can use np.array_split which deals with uneven splits quite superbly. First, I'll initialise some random array:

arr = np.random.randn(87, 40)

Next, compute the indices to split on. If the shape of arr is divisible by 40, then generate even splits. Otherwise, overflows go into the (n - 1)th array.

# compute the indices to split on
if arr.shape[0] % 40 == 0:
    split_idx = arr.shape[0] // 40
else:
    split_idx = np.arange(40, arr.shape[0], 40)[:-1]

Finally, call array_split, and split on split_idx:

# split along the 0th axis
splits = np.array_split(arr, split_idx, axis=0)

Verify that our array was partitioned correctly:

[s.shape for s in splits]
[(40, 40), (47, 40)]

Upvotes: 6

Related Questions