Bipartite graphs

#!/usr/bin/env python

from collections import defaultdict


def bipartite(*pairs):
    """Construct a bipartition of nodes

    The input is a series of pairs of nodes that must be in different disjoint sets.

    The idea is this:
    1. Construct a graph where edges join nodes that must be distinct.
    2. Pick a node that's not yet been visited. Put it into an arbitrary side.
    3. For all nodes connected to the node, put them into the other side.
    4. Repeat until we've allocated all nodes that are transitively linked to the first.
    5. If there are any nodes left, go to step two and pick another.

    During steps 3&4, it may be the case that a node would be allocated to both sides.
    That indicates that there is some path of odd length from the node to itself."""

    # Process the pairs into { node: { linked nodes } }
    nodes = defaultdict(set)
    for a, b in pairs:
        nodes[a].add(b)
        nodes[b].add(a)

    sides = (set(), set())
    side = 0
    while len(nodes) > 0:
        node, adjoints = nodes.popitem()
        if node in sides[1-side]:
            return None
        sides[side].add(node)
        while len(adjoints) > 0:
            side = 1 - side
            processing = adjoints
            adjoints = set()  # Work out the next tranche to handle
            for node in processing:
                if node in sides[1 - side]:
                    return None
                sides[side].add(node)
                if node in nodes:   # If we've not already handled this
                    adjoints.update(nodes.pop(node))

    return set(frozenset(s) for s in sides)


def test_bipartite():
    assert bipartite(('a', 'b')) == {frozenset({'a'}), frozenset({'b'})}
    assert bipartite(('a', 'b'), ('b', 'c'), ('c', 'a')) is None
    assert bipartite(('a', 'b'), ('b', 'c'), ('a', 'd'), ('c', 'd')) == {frozenset({'a', 'c'}), frozenset({'b', 'd'})}