import pandas as pd

# 假设这是你的DataFrame
df = pd.DataFrame({
    'from': ['A', 'B', 'C', 'A', 'D', 'E', 'G', 'H', 'A', 'A'],
    'to': ['B', 'C', 'A', 'D', 'E', 'H', 'I', 'J', 'F', 'G']
})

def get_connected_edges(df, start_node, max_connections=6):
    # 初始化结果列表
    result = []
    # 初始化待检查的节点集合，并设置起始节点的深度为0
    to_check = {start_node: 0}
    # 初始化已检查的节点集合
    checked = set()
    # 初始化连接数字典
    connection_counts = {node: 0 for node in df['from'].tolist() + df['to'].tolist()}
    
    while to_check:
        # 取出一个节点及其深度
        node, depth = to_check.popitem()
        # 添加到已检查集合
        checked.add(node)
        
        # 如果节点连接数超过最大限制，则跳过
        if connection_counts[node] > max_connections:
            continue
        
        # 找到所有从这个节点出发的边
        from_edges = df[df['from'] == node]
        # 找到所有进入这个节点的边
        to_edges = df[df['to'] == node]
        
        # 更新连接数
        connection_counts[node] += len(from_edges) + len(to_edges)
        
        # 如果连接数超过最大限制，则跳过
        if connection_counts[node] > max_connections:
            continue
        
        # 将这些边添加到结果列表，并记录深度
        for index in from_edges.index.tolist() + to_edges.index.tolist():
            result.append({'index': index, 'depth': depth})
        
        # 将所有连接的节点添加到待检查集合，并更新深度
        for edge in from_edges['to'].tolist() + to_edges['from'].tolist():
            if edge not in checked:
                to_check[edge] = depth + 1
    
    # 将结果列表转换为DataFrame
    result_df = pd.DataFrame(result)
    print(result_df)
    # 使用找到的索引来筛选DataFrame，并添加深度列
    connected_edges_df = df.loc[result_df['index']].assign(depth=result_df['depth'])
    print(connected_edges_df)
    
    return connected_edges_df

# 假设我们要找到与节点 'A' 相关的所有边
start_node = 'A'
connected_edges_df = get_connected_edges(df, start_node)

# print(connected_edges_df)
