Source code for networkx.algorithms.community.kernighan_lin

"""Functions for computing the Kernighan–Lin bipartition algorithm."""

from itertools import count

import networkx as nx
from networkx.algorithms.community.community_utils import is_partition
from networkx.utils import BinaryHeap, not_implemented_for, py_random_state

__all__ = ["kernighan_lin_bisection"]


def _kernighan_lin_sweep(edges, side):
    """
    This is a modified form of Kernighan-Lin, which moves single nodes at a
    time, alternating between sides to keep the bisection balanced.  We keep
    two min-heaps of swap costs to make optimal-next-move selection fast.
    """
    costs0, costs1 = costs = BinaryHeap(), BinaryHeap()
    for u, side_u, edges_u in zip(count(), side, edges):
        cost_u = sum(w if side[v] else -w for v, w in edges_u)
        costs[side_u].insert(u, cost_u if side_u else -cost_u)

    def _update_costs(costs_x, x):
        for y, w in edges[x]:
            costs_y = costs[side[y]]
            cost_y = costs_y.get(y)
            if cost_y is not None:
                cost_y += 2 * (-w if costs_x is costs_y else w)
                costs_y.insert(y, cost_y, True)

    i = 0
    totcost = 0
    while costs0 and costs1:
        u, cost_u = costs0.pop()
        _update_costs(costs0, u)
        v, cost_v = costs1.pop()
        _update_costs(costs1, v)
        totcost += cost_u + cost_v
        i += 1
        yield totcost, i, (u, v)


[docs] @not_implemented_for("directed") @py_random_state(4) @nx._dispatchable(edge_attrs="weight") def kernighan_lin_bisection(G, partition=None, max_iter=10, weight="weight", seed=None): """Partition a graph into two blocks using the Kernighan–Lin algorithm. This algorithm partitions a network into two sets by iteratively swapping pairs of nodes to reduce the edge cut between the two sets. The pairs are chosen according to a modified form of Kernighan-Lin [1]_, which moves node individually, alternating between sides to keep the bisection balanced. Parameters ---------- G : NetworkX graph Graph must be undirected. partition : tuple Pair of iterables containing an initial partition. If not specified, a random balanced partition is used. max_iter : int Maximum number of times to attempt swaps to find an improvement before giving up. weight : key Edge data key to use as weight. If None, the weights are all set to one. seed : integer, random_state, or None (default) Indicator of random number generation state. See :ref:`Randomness<randomness>`. Only used if partition is None Returns ------- partition : tuple A pair of sets of nodes representing the bipartition. Raises ------ NetworkXError If partition is not a valid partition of the nodes of the graph. References ---------- .. [1] Kernighan, B. W.; Lin, Shen (1970). "An efficient heuristic procedure for partitioning graphs." *Bell Systems Technical Journal* 49: 291--307. Oxford University Press 2011. """ n = len(G) labels = list(G) seed.shuffle(labels) index = {v: i for i, v in enumerate(labels)} if partition is None: side = [0] * (n // 2) + [1] * ((n + 1) // 2) else: try: A, B = partition except (TypeError, ValueError) as err: raise nx.NetworkXError("partition must be two sets") from err if not is_partition(G, (A, B)): raise nx.NetworkXError("partition invalid") side = [0] * n for a in A: side[index[a]] = 1 if G.is_multigraph(): edges = [ [ (index[u], sum(e.get(weight, 1) for e in d.values())) for u, d in G[v].items() ] for v in labels ] else: edges = [ [(index[u], e.get(weight, 1)) for u, e in G[v].items()] for v in labels ] for i in range(max_iter): costs = list(_kernighan_lin_sweep(edges, side)) min_cost, min_i, _ = min(costs) if min_cost >= 0: break for _, _, (u, v) in costs[:min_i]: side[u] = 1 side[v] = 0 A = {u for u, s in zip(labels, side) if s == 0} B = {u for u, s in zip(labels, side) if s == 1} return A, B