a02_a_is_bipartite_graph.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Check if the graph is bipartite.
"""

import collections


class Graph(object):
    """
    Simple graph model.
    """

    def __init__(self, num_nodes):
        self.edges = collections.defaultdict(dict)
        self.num_nodes = num_nodes

    def make_edge(self, from_, to, value):
        self.edges[from_][to] = value


def is_bipartite(graph):
    """
    Returns True if the graph is bipartite.
    """

    def bfs(v):
        q = [(v, 1)]

        while q:
            x, val = q.pop(0)
            if x in memo:
                continue
            memo[x] = val
            ts = graph.edges.get(x, {}).keys()
            if any(memo.get(a) == val for a in ts):
                return False
            q += [(a, -val) for a in ts]
        return True

    memo = {}
    return all(w in memo or bfs(w) for w in xrange(graph.num_nodes))


if __name__ == '__main__':
    a = Graph(4)
    a.make_edge(0, 1, 1)
    a.make_edge(1, 2, 1)
    a.make_edge(2, 3, 1)
    a.make_edge(3, 0, 1)
    #
    #    0 - 1
    #    |   |
    #    3 - 2
    #

    b = Graph(5)
    b.make_edge(0, 1, 1)
    b.make_edge(0, 2, 1)
    b.make_edge(0, 3, 1)
    b.make_edge(0, 4, 1)
    #
    #    1   2
    #     \ /
    #      0
    #     / \
    #    4   3
    #

    c = Graph(3)
    c.make_edge(0, 1, 1)
    c.make_edge(1, 2, 1)
    c.make_edge(2, 0, 1)
    #
    #
    #    0 - 1
    #     \ /
    #      2
    #

    print([is_bipartite(g) for g in (a, b, c)])