How would you split a numpy array where the elements give the partition size?

user_12960856 0
Viewed 66

For an numpy 1d array such as:

In [1]: A = np.array([2,5,1,3,9,0,7,4,1,2,0,11])
In [2]: A
Out[2]: array([2,5,1,3,9,0,7,4,1,2,0,11])

I need to split the array by using the values as a sub-array length.

For the example array:

  1. The first index has a value of 2, so I need the first split to occur at index 0 + 2, so it would result in ([2,5,1]).
  2. Skip to index 3 (since indices 0-2 were gobbled up in step 1).
  3. The value at index 3 = 3, so the second split would occur at index 3 + 3, and result in ([3,9,0,7]).
  4. Skip to index 7
  5. The value at index 7 = 4, so the third and final split would occur at index 7 + 4, and result in ([4,1,2,0,11])

I'm using this simple array as an example, because I think it will help in my actual use case, which is reading data from binary files (either as bytes or unsigned shorts). I'm guessing that numpy will be the fastest way to do it, but I could also use struct/bytearray/lists or whatever would be best.

I hope this makes sense. I had a hard time trying to figure out how best to word the question.

numpy arrays split
1 Answers

Here is an approach using standard python lists and a while loop:

def custom_partition(arr):
    partitions = []
    i = 0
    while i < len(arr):
        pariton_size = arr[i]
        next_i = i + pariton_size + 1
        partitions.append(arr[i:next_i])
        i = next_i
    return partitions

a = [2, 5, 1, 3, 9, 0, 7, 4, 1, 2, 0, 11]
b = custom_partition(a)
print(b)

Output:

[[2, 5, 1], [3, 9, 0, 7], [4, 1, 2, 0, 11]]
user_63282560
Related