Spaces:
Running
Running
| import random | |
| from collections.abc import Iterable | |
| from typing import Any | |
| from graphgen.bases import BaseGraphStorage, BasePartitioner | |
| from graphgen.bases.datatypes import Community | |
| NODE_UNIT: str = "n" | |
| EDGE_UNIT: str = "e" | |
| class DFSPartitioner(BasePartitioner): | |
| """ | |
| DFS partitioner that partitions the graph into communities of a fixed size. | |
| 1. Randomly choose a unit. | |
| 2. Random walk using DFS until the community reaches the max unit size. | |
| (In GraphGen, a unit is defined as a node or an edge.) | |
| """ | |
| def partition( | |
| self, | |
| g: BaseGraphStorage, | |
| max_units_per_community: int = 1, | |
| **kwargs: Any, | |
| ) -> Iterable[Community]: | |
| nodes = g.get_all_nodes() | |
| edges = g.get_all_edges() | |
| adj, _ = self._build_adjacency_list(nodes, edges) | |
| used_n: set[str] = set() | |
| used_e: set[frozenset[str]] = set() | |
| units = [(NODE_UNIT, n[0]) for n in nodes] + [ | |
| (EDGE_UNIT, frozenset((u, v))) for u, v, _ in edges | |
| ] | |
| random.shuffle(units) | |
| for kind, seed in units: | |
| if (kind == NODE_UNIT and seed in used_n) or ( | |
| kind == EDGE_UNIT and seed in used_e | |
| ): | |
| continue | |
| comm_n, comm_e = [], [] | |
| stack = [(kind, seed)] | |
| cnt = 0 | |
| while stack and cnt < max_units_per_community: | |
| k, it = stack.pop() | |
| if k == NODE_UNIT: | |
| if it in used_n: | |
| continue | |
| used_n.add(it) | |
| comm_n.append(it) | |
| cnt += 1 | |
| for nei in adj[it]: | |
| e_key = frozenset((it, nei)) | |
| if e_key not in used_e: | |
| stack.append((EDGE_UNIT, e_key)) | |
| break | |
| else: | |
| if it in used_e: | |
| continue | |
| used_e.add(it) | |
| comm_e.append(tuple(it)) | |
| cnt += 1 | |
| # push neighboring nodes | |
| for n in it: | |
| if n not in used_n: | |
| stack.append((NODE_UNIT, n)) | |
| if comm_n or comm_e: | |
| yield Community(id=seed, nodes=comm_n, edges=comm_e) | |