Python循环优化

Question

假设有以下这样的DataFrame：

     id  days  cluster
0   aaa     0        0
1   bbb     0        0
2   ccc     0        1
3   ddd     0        1
4   eee     0        0
5   fff     0        1
6   ggg     1        0
7   hhh     1        1
8   iii     1        0
9   lll     1        1
10  mmm     1        1
11  aaa     1        3
12  bbb     1        3

我的目标是创建一个字典，其中键是元素为“id”列的元组，而值是“cluster”列的元素列表，如果两个“id”具有相同的“cluster”值，则对所有过滤后的“days”列进行筛选。即，如果“days”发生变化，但存在具有相同“cluster”值的“id”元组，则我希望将此值添加到我的已有列表中。所需输出如下所示：

{('aaa', 'bbb'): [0, 3],('aaa', 'eee'): [0], ('bbb', 'eee'): [0], ('ccc', 'ddd'): [1], 
('ccc', 'fff'): [1], ('ddd', 'fff'): [1], ('ggg', 'iii'): [0],
 ('hhh', 'lll'): [1], ('hhh', 'mmm'): [1], ('lll', 'mmm'): [1]}

我使用以下代码片段得出了这个结果，但是当数据行数达到百万级别时，速度变得太慢了。如何优化代码？

y={}
for i in range(0, max(df.iloc[:,1]) + 1):
    x = df.loc[df['days'] == i]
    for j in range(0,l en(x)):
        for z in range(1, len(x)):
            if (x.iloc[z,0], x.iloc[j,0]) in y:
                pass
            else:
             if (x.iloc[j,0], x.iloc[z,0]) not in y:
                 if x.iloc[j,0] != x.iloc[z,0] and x.iloc[j,2] == x.iloc[z,2]:
                     y[(x.iloc[j,0], x.iloc[z,0])] = [x.iloc[j,2]]
             else:
                 if x.iloc[j,0] != x.iloc[z,0] and x.iloc[j,2] == x.iloc[z,2]:
                     y[(x.iloc[j,0], x.iloc[z,0])].append(x.iloc[j,2])

- Nicolapo

---

在你的例子中，ID 'aaa' 可能的簇值为 0 和 3（分别对应第 0 天和第 1 天）。但是在你期望的输出中，ID 'aaa' 与 'ccc'、'ddd'、'fff'、'hhh'、'lll' 和 'mmm' 分组，它们的簇值为 1 或 2。因此我不理解你的语句 if the two 'id' have the same 'cluster' value。 - mtrw

@mtrw 你是对的！已经修复了，我之前发布的期望输出是错误的！谢谢。 - Nicolapo

Answer 1

考虑到瓶颈在于获取id的组合，为什么不把它留到最后呢？

按id分组数据，每个id将显示其所在的“bin”（天数、集群）的集合：

grouped = collections.defaultdict(set)
for index, (id_, day, cluster) in df.iterrows():
    grouped[id_].add((day, cluster))

对于找到的每个垃圾桶组合，请列出属于每个垃圾桶组合的ID列表：

binned = collections.defaultdict(list)
for id_, bins in grouped.items():
    binned[tuple(sorted(bins))].append(id_)

只有当你需要时，才通过聚类来简化：

clustered = collections.defaultdict(list)
for bins, ids in binned.items():
    clusters = set(cluster for (day, cluster) in bins)
    clustered[tuple(sorted(clusters))].extend(ids)

最后，获取每个聚类箱的id组合不应该是问题：

for bins, ids in clustered.items():
    if len(ids) > 1:
        for comb_id in itertools.combinations(ids, 2):
            print(bins, comb_id) 
            # or do other stuff with it

- chapelo

Answer 2

你可以利用 pandas.DataFrame.groupby 方法：

result = collections.defaultdict(list)

for (day, cluster), group in df.groupby(["days", "cluster"]):
    for comb in itertools.combinations(df["id"][group.index], 2):
        result[comb].append(cluster)

这将为您提供所需的结果：

defaultdict(<class 'list'>, {('aaa', 'bbb'): [0, 3], ('aaa', 'eee'): [0], ('bbb', 'eee'): [0], ('ccc', 'ddd'): [1], ('ccc', 'fff'): [1], ('ddd', 'fff'): [1], ('ggg', 'iii'): [0], ('hhh', 'lll'): [1], ('hhh', 'mmm'): [1], ('lll', 'mmm'): [1]})

- chapelo

---

你的代码看起来更快，但不幸的是，在使用了所有可用的内存后，我的会话在处理200,000行、9,000个不同的“id”、2个不同的“days”和8个不同的“cluster”时崩溃了。 - Nicolapo

这并不奇怪。如果数字是随机的，你会预计在每个簇中看到几乎所有的ID出现多次。组合是“无限”的。你尝试过将簇组合作为字典键，将出现在该簇组合中的ID列表作为字典值的反向方法吗？然后获取每个簇组合的ID组合就变得可行了。我使用你提到的数字的随机样本，在使用6GB RAM的约2分钟内完成了它。你也可以尝试生产者-消费者方法。 - chapelo