Split numpy array into groups of the same number

Question

I have an array :

A = np.array([1,2,2,2,3,3,1,1,1,4,4,4,4,4,5,5])

I want to do some function that will split A into smaller arrays (or lists, either works - I will use lists in the example easier to write), each made up of a chunk of the same number.

B = stackoverflowgeniusfunc(A)
B
[[1],[2,2,2],[3,3],[1,1,1],[4,4,4,4,4],[5,5]]

How do I do this?

Does this answer your question? [How do I use itertools.groupby()?](https://stackoverflow.com/questions/773/how-do-i-use-itertools-groupby) — B Remmelzwaal, May 23 '23 at 14:43

score 1 · Answer 1 · answered May 23 '23 at 14:51

Assuming that you want to group the similar objects together, you can do something like this.

    def stackoverflowgeniusfunc(arr):
    result = []
    current_chunk = [arr[0]]

    for i in range(1, len(arr)):
        if arr[i] == arr[i-1]:
            current_chunk.append(arr[i])
        else:
            result.append(current_chunk)
            current_chunk = [arr[i]]

    result.append(current_chunk)
    return result

The function splits the array A into smaller arrays/lists, where each chunk consists of consecutive elements of the same number.

PS: You can also use numpy's

np.split()

function with a custom condition to split the array whenever the value changes

    import numpy as np
    
    A = np.array([1, 2, 2, 2, 3, 3, 1, 1, 1, 4, 4, 4, 4, 4, 5, 5])
    
    # Find the indices where the value changes
    change_indices = np.where(A[1:] != A[:-1])[0] + 1
    
    # Split the array based on the change indices
    B = np.split(A, change_indices)

Talha Tayyab · Answer 2 · 2023-05-23T16:54:34.303

1

from itertools import groupby
import numpy as np
A=np.array([1,2,2,2,3,3,1,1,1,4,4,4,4,4,5,5])

def stackoverflowgeniusfunc(a):
  return [list(group[1]) for group in groupby(A)]

stackoverflowgeniusfunc(A)

#output
[[1], [2, 2, 2], [3, 3], [1, 1, 1], [4, 4, 4, 4, 4], [5, 5]]

edited May 23 '23 at 16:54

answered May 23 '23 at 15:02

Talha Tayyab

8,111
25
27
44

score 1 · Accepted Answer · answered May 23 '23 at 16:45

1

numpy does not support variable size dimensions so your expected output cannot be a numpy array. It can be a normal Python list however. groupby would be the usual way to obtain it:

from itertools import groupby

A = [1,2,2,2,3,3,1,1,1,4,4,4,4,4,5,5]
G = [g for _,(*g,) in groupby(A)]

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

Note: this would also work if A were a numpy array but there is no point in using numpy for this unless you need is elsewhere

answered May 23 '23 at 16:45

Alain T.

40,517
4
31
51

Sorry for the delay. Simple and perfect thank you! You are correct that A is a numpy array for use elsewhere, the compatibility with arrays and lists is neat though. Thanks! – Jack Jun 01 '23 at 09:10

Split numpy array into groups of the same number

3 Answers3