在python中基于另一个数组上的重复值合并数组？

Question

我已经将我的数据组织成3个列表。第一个简单地包含浮点数，其中一些是重复的。第二个和第三个列表包含长度可变的1D数组

第一个列表被排序，所有列表包含相同数量的元素

总体格式如下：

a = [1.0, 1.5, 1.5, 2 , 2]
b = [arr([1 2 3 4 10]), arr([4 8 10 11 5 6 12]), arr([1 5 7]), arr([70 1 2]), arr([1])]
c = [arr([3 4 8]), arr([5 6 12]), arr([6 7 10 123 14]), arr([70 1 2]), arr([1 5 10 4])]

我试图找到一种方法来合并列表b和c中的数组，如果它们对应的浮点数在列表a中相同。对于上述示例，预期结果将是：

a = [1.0, 1.5, 2]
b = [arr([1 2 3 4 10]), arr([4 8 10 11 5 6 12 1 5 7]), arr([70 1 2 1])]
c = [arr([3 4 8]), arr([5 6 12 6 7 10 123 14]), arr([70 1 2 1 5 10 4]])]

我该怎么做呢？这和拉链有关吗

Answer 1

由于a已排序，因此我将使用itertools.groupby。与@MadPhysicast的答案类似，但在zip列表上迭代：

import numpy as np
from itertools import groupby

arr = np.array

a = [1.0, 1.5, 1.5, 2 , 2]
b = [arr([1, 2, 3, 4, 10]), arr([4, 8, 10, 11, 5, 6, 12]), arr([1, 5, 7]), arr([70, 1, 2]), arr([1])]
c = [arr([3, 4, 8]), arr([5, 6, 12]), arr([6, 7, 10, 123, 14]), arr([70, 1, 2]), arr([1, 5, 10, 4])]

res_a, res_b, res_c = [], [], []
for k, g in groupby(zip(a, b, c), key=lambda x: x[0]):
    g = list(g)
    res_a.append(k)
    res_b.append(np.concatenate([x[1] for x in g]))
    res_c.append(np.concatenate([x[2] for x in g]))

…输出res_a、res_b和res_c为：

[1.0, 1.5, 2]
[array([ 1,  2,  3,  4, 10]), array([ 4,  8, 10, 11,  5,  6, 12,  1,  5,  7]), array([70,  1,  2,  1])]
[array([3, 4, 8]), array([  5,   6,  12,   6,   7,  10, 123,  14]), array([70,  1,  2,  1,  5, 10,  4])]

或者，如果a未排序，您可以使用defaultdict：

import numpy as np
from collections import defaultdict

arr = np.array

a = [1.0, 1.5, 1.5, 2 , 2]
b = [arr([1, 2, 3, 4, 10]), arr([4, 8, 10, 11, 5, 6, 12]), arr([1, 5, 7]), arr([70, 1, 2]), arr([1])]
c = [arr([3, 4, 8]), arr([5, 6, 12]), arr([6, 7, 10, 123, 14]), arr([70, 1, 2]), arr([1, 5, 10, 4])]

res_a, res_b, res_c = [], [], []

d = defaultdict(list)

for x, y, z in zip(a, b, c):
    d[x].append([y, z])

for k, v in d.items():
    res_a.append(k)
    res_b.append(np.concatenate([x[0] for x in v]))
    res_c.append(np.concatenate([x[1] for x in v]))