根据数组中的唯一值拆分数组

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”中存在三个唯一值的情况。如果数组“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]]