fix: add special edge handling in bfs, dfs and base partitioner (#150)

* fix: add special edge handling in bfs, dfs and base partitioner

* refactor: extract a helper method in the base class

* fix: add List import to dfs_partitioner

* fix: align edge type

* fix: fix lint error

---------

Co-authored-by: chenzihong <522023320011@smail.nju.edu.cn>

Author: CHERRY-ui8

graphgen/bases/base_partitioner.py

@@ -39,14 +39,17 @@ def community2batch(
             if node_data:
                 nodes_data.append((node, node_data))
         edges_data = []
-        for u, v in edges:
-            edge_data = g.get_edge(u, v)
+        for edge in edges:
+            # Filter out self-loops and invalid edges
+            if not isinstance(edge, tuple) or len(edge) != 2:
+                continue
+            u, v = edge
+            if u == v:
+                continue
+
+            edge_data = g.get_edge(u, v) or g.get_edge(v, u)
             if edge_data:
                 edges_data.append((u, v, edge_data))
-            else:
-                edge_data = g.get_edge(v, u)
-                if edge_data:
-                    edges_data.append((v, u, edge_data))
         return nodes_data, edges_data
 
     @staticmethod
@@ -61,9 +64,11 @@ def _build_adjacency_list(
         """
         adj: dict[str, List[str]] = {n[0]: [] for n in nodes}
         edge_set: set[tuple[str, str]] = set()
-        for e in edges:
-            adj[e[0]].append(e[1])
-            adj[e[1]].append(e[0])
-            edge_set.add((e[0], e[1]))
-            edge_set.add((e[1], e[0]))
+        for u, v, _ in edges:
+            if u == v:
+                continue
+            adj[u].append(v)
+            adj[v].append(u)
+            edge_set.add((u, v))
+            edge_set.add((v, u))
         return adj, edge_set

graphgen/models/partitioner/bfs_partitioner.py

@@ -63,9 +63,7 @@ def partition(
                     if it in used_e:
                         continue
                     used_e.add(it)
-
-                    u, v = it
-                    comm_e.append((u, v))
+                    comm_e.append(tuple(sorted(it)))
                     cnt += 1
                     # push nodes that are not visited
                     for n in it:

graphgen/models/partitioner/dfs_partitioner.py

@@ -1,6 +1,6 @@
 import random
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, List
 
 from graphgen.bases import BaseGraphStorage, BasePartitioner
 from graphgen.bases.datatypes import Community
@@ -42,7 +42,8 @@ def partition(
             ):
                 continue
 
-            comm_n, comm_e = [], []
+            comm_n: List[str] = []
+            comm_e: List[tuple[str, str]] = []
             stack = [(kind, seed)]
             cnt = 0
 
@@ -63,7 +64,7 @@ def partition(
                     if it in used_e:
                         continue
                     used_e.add(it)
-                    comm_e.append(tuple(it))
+                    comm_e.append(tuple(sorted(it)))
                     cnt += 1
                     # push neighboring nodes
                     for n in it:

graphgen/models/partitioner/ece_partitioner.py

@@ -142,7 +142,7 @@ def _add_unit(u):
             return Community(
                 id=seed_unit[1],
                 nodes=list(community_nodes.keys()),
-                edges=[tuple(edge) for edge in community_edges if isinstance(edge, frozenset) and len(edge)==2],
+                edges=[tuple(sorted(e)) for e in community_edges]
             )
 
         for unit in tqdm(all_units, desc="ECE partition"):

graphgen/utils/help_nltk.py

@@ -3,13 +3,14 @@
 from typing import Dict, List, Final, Optional
 import warnings
 import nltk
-import jieba
-
 warnings.filterwarnings(
     "ignore",
     category=UserWarning,
     module=r"jieba\._compat"
 )
+# pylint: disable=wrong-import-position
+import jieba
+
 
 class NLTKHelper:
     """