Разбиение массивов в зависимости от уникальных значений в массиве

Question

В настоящее время у меня есть два массива, один из которых имеет несколько повторяющихся значений, а другой с уникальными значениями.

Например, массив 1: a = [1, 1, 2, 2, 3, 3]

Например, массив 2: b = [10, 11, 12, 13, 14, 15]

Я разрабатывал код на python, который просматривает первый массив и различает элементы, которые являются одинаковыми, и запоминает индексы. Создается новый массив, который содержит элементы массива b по этим индексам.

Например: так как массив 'a' имеет три уникальных значения в позициях 1,2 ... 3,4 ... 5,6, то будут созданы три новых массива так, чтобы он содержал элементы массива b в позициях 1,2 ... 3,4 ... 5,6. Таким образом, результатом будет три новых массива:

b1 = [10, 11]

b2 = [12, 13]

b3 = [14, 15]

Мне удалось разработать код, однако он работает только для трех уникальных значений в массиве «a». Если в массиве «а» есть более или менее уникальные значения, код должен быть физически изменен.

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)

Мне было интересно, есть ли, возможно, более элегантный способ выполнить эту задачу, чтобы код мог справиться с n числом уникальных значений.

Answer 1

Хорошо, учитывая, что вы также используете numpy, вот один из способов использования np.unique. Вы можете установить return_index=True, чтобы получить индексы уникальных значений, и использовать их для разделения массива b на 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:])

Выход

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

Answer 2

Вы можете использовать функцию 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]]