如何檢查 python 中的數據組是否連接

Question

假設我們有一個這樣的數據集：

import pandas as pd
import numpy as np

df = pd.DataFrame({'ID':[1,1,1,2,2,2,3,3,3,4,4,4], 'Number':[234, 43, 455, 112, 45, 234, 982, 41, 112, 46, 109, 4]})

現在，如果您運行上面的代碼，您會看到 ID 的 1,2 和 3 都是“連接”的，即 ID 1 通過 Number == 234 連接到 ID 2，ID 2 通過 Number 連接到 ID 3 == 112。但是 ID 4 沒有連接到其他 3 個 ID。 所以我想要一個像這樣的 df_final

df_final = pd.DataFrame({'ID':[1,1,1,2,2,2,3,3,3,4,4,4], 'Number':[234, 43, 455, 112, 45, 234, 982, 41, 112, 46, 109, 4], 'Group':[1,1,1,1,1,1,1,1,1,2,2,2]})

我的數據集很大，所以我需要一些盡可能高效的代碼。 我定義了一個 'is_a_neighbor' function 來檢查兩個 ID 之間是否有重疊。 如果連接了兩個 ID（即有重疊），則返回“True”。 我想我需要以某種方式將它嵌入到一個循環中。 這是 function：

def is_a_neighbor(id_1, id_2):
  val = np.isin(df[df['ID'] == id_1][['Number']].to_numpy(), df[df['ID'] == id_2][['Number']].to_numpy()).any()
  return val

和試運行

# Test is_a_neighbor function
id_1 = 1
id_2 = 2

val = is_a_neighbor(id_1, id_2)
val

這將返回“真”。 感謝您對此提供任何幫助/指導。

Answer 1

以下代碼應該可以工作：

import pandas as pd
import networkx as nx

df = pd.DataFrame({'ID':[1,1,1,2,2,2,3,3,3,4,4,4], 
                   'Number':[234, 43, 455, 112, 45, 234, 982, 41, 112, 46, 109, 4]})

df.ID = df.ID.astype(str)
g = nx.from_pandas_edgelist(df,'ID','Number')
connected_components = nx.connected_components(g)
d = {y:i for i, x in enumerate(connected_components) for y in x}

df['group'] = df.ID.replace(d)
df.ID = df.ID.astype(int)

我們將ID列轉換為字符串的原因是因為我們希望ID列中的1與Number列中的1不同。

如果我們不這樣做，以下示例將僅顯示一組：

ID	數字
1	3
2	3
3	4

為了擴展代碼的速度，以下測試的平均時間為 3.4 秒：

import pandas as pd
import numpy as np
import networkx as nx
import timeit

test_df = pd.DataFrame({'ID': np.random.randint(0, 10000, 10000), 
                        'Number': np.random.randint(0, 10000, 10000)})

def add_group(df):
    df.ID = df.ID.astype(str)
    g = nx.from_pandas_edgelist(df,'ID','Number')
    connected_components = nx.connected_components(g)
    d = {y:i for i, x in enumerate(connected_components) for y in x}

    df['group'] = df.ID.replace(d)
    df.ID = df.ID.astype(int)
    return df

timeit.timeit(lambda : add_group(test_df), number=100)

Answer 2

據我所知，如果有任何數字 N 使得 (i1, N) 和 (id2, N) 都存在，那么您有一長串對 (Id, Number) 和兩個不同的 Id 屬於同一個組。

我不知道為什么您認為為此需要 panda 或 numpy

為什么不這樣：

data = [(1, 234), (1, 43), (1, 455), (1, 112), (2, 234), ...]


# Create sets of numbers for each index

number_sets = {}
for (index, number) in data:
    # if we've not seen the index before we create an empty set
    if index not in number_sets:
        number_sets[index] = set()

    number_sets[index].add(number)

現在我們創建一個組列表，對於每個組，我們記錄一個索引列表和使用這些索引中的任何一個找到的一組數字

groups = []
    
for (index, numbers) in number_sets.items():
    for (index_list, superset) in groups:
        # if any of the numbers are in the superset then this index belongs to this group
        if not superset.isdisjoint(numbers):
           index_list.append(index)
           superset.join(numbers)
           break
    else:
        # We never had a break so we've found a new group
        groups.append(([index], numbers))