CODEWITHSUNDEEP
pythonarrayspython-3.xsorting
JasKa071
JasKa071

Reputation: 55

Splitting arrays depending on unique values in an array

I currently have two arrays, one of which has several repeated values and another with unique values.

Eg array 1 : a = [1, 1, 2, 2, 3, 3]

Eg array 2 : b = [10, 11, 12, 13, 14, 15]

I was developing a code in python that looks at the first array and distinguishes the elements that are all the same and remembers the indices. A new array is created that contains the elements of array b at those indices.

Eg: As array 'a' has three unique values at positions 1,2... 3,4... 5,6, then three new arrays would be created such that it contains the elements of array b at positions 1,2... 3,4... 5,6. Thus, the result would be three new arrays:

b1 = [10, 11]

b2 = [12, 13]

b3 = [14, 15]

I have managed to develop a code, however, it only works for when there are three unique values in array 'a'. In the case there are more or less unique values in array 'a', the code has to be physically modified.


import itertools
import numpy as np
import matplotlib.tri as tri
import sys

a = [1, 1, 2, 2, 3, 3]
b = [10, 10, 20, 20, 30, 30]
b_1 = []
b_2 = []
b_3 = []



unique = []
for vals in a:
    if vals not in unique:
        unique.append(vals)

if len(unique) != 3:
    sys.exit("More than 3 'a' values  - check dimension")

for j in range(0,len(a)):
    if a[j] == unique[0]:
        b_1.append(c[j])
    elif a[j] == unique[1]:
        b_2.append(c[j])
    elif a[j] == unique[2]:
        b_3.append(c[j])
    else:
        sys.exit("More than 3 'a' values  - check dimension")
print (b_1)
print (b_2)
print (b_3)

I was wondering if there is perhaps a more elegant way to perform this task such that the code is able to cope with an n number of unique values.

Upvotes: 3

Views: 2791

Answers (2)

Mykola Zotko
Mykola Zotko

Reputation: 17794

You can use the function groupby():

from itertools import groupby
from operator import itemgetter

a = [1, 1, 2, 2, 3, 3]
b = [10, 11, 12, 13, 14, 15]

[[i[1] for i in g] for _, g in groupby(zip(a, b), key=itemgetter(0))]
# [[10, 11], [12, 13], [14, 15]]

Upvotes: 2

yatu
yatu

Reputation: 88226

Well given that you are also using numpy, here's one way using np.unique. You can set return_index=True to get the indices of the unique values, and use them to split the array b with np.split:

a = np.array([1, 1, 2, 2, 3, 3])
b = np.array([10, 11, 12, 13, 14, 15])

u, s = np.unique(a, return_index=True)
np.split(b,s[1:])

Output 

[array([10, 11]), array([12, 13]), array([14, 15])]

Upvotes: 8

Related Questions