dsc40graph.py

"""Efficient undirected and directed graph data structures."""


MAX_NODES_DISPLAYED = 20
MAX_EDGES_DISPLAYED = 10


class Error(Exception):
    """Base class for all errors specific to this module."""


class DoesNotExistError(Error):
    """The node/edge does not exist."""


def _take(iterable, n):
    for i, x in enumerate(iterable):
        if i >= n:
            break
        yield x


class _NodesView:
    """A view into a graph's nodes."""

    def __init__(self, nodes):
        self._nodes = nodes

    def __contains__(self, node):
        """Membership query."""
        return node in self._nodes

    def __len__(self):
        return len(self._nodes)

    def __iter__(self):
        yield from self._nodes

    def __repr__(self):
        limit = MAX_NODES_DISPLAYED
        over_limit = len(self._nodes) > limit

        if over_limit:
            nodes_to_print = list(_take(self._nodes, MAX_NODES_DISPLAYED))
            suffix = "..."
        else:
            nodes_to_print = self._nodes
            suffix = ""

        node_string = ", ".join(repr(node) for node in nodes_to_print) + suffix
        return f"<{len(self._nodes)} nodes: {node_string}>"


class _EdgeView:
    """Base class for views into a graph's edges."""

    def __init__(self, adj, number_of_edges):
        self._adj = adj
        self._number_of_edges = number_of_edges

    def __contains__(self, edge):
        """Perform an edge query.
        
        Average case time complexity: Theta(1)
        """
        u, v = edge
        try:
            return v in self._adj[u]
        except KeyError:
            return False

    def __len__(self):
        """The number of edges.

        Average case time complexity: Theta(1)
        """
        return self._number_of_edges

    def __repr__(self):
        limit = MAX_EDGES_DISPLAYED
        over_limit = len(self) > limit

        if over_limit:
            edges_to_print = list(_take(self, MAX_EDGES_DISPLAYED))
            suffix = "..."
        else:
            edges_to_print = list(self)
            suffix = ""

        edges_string = ", ".join(repr(edge) for edge in edges_to_print) + suffix
        return f"<{len(self)} edges: {edges_string}>"


class _UndirectedEdgeView(_EdgeView):
    """A view into an undirected graph's edges."""

    def __iter__(self):
        """Iterate through the edges.

        Each edge in the graph is yielded exactly once as a pair whose order is
        arbitrary. That is, suppose that a graph has an edge between node 1 and
        node 2. Then the pair (1,2) may be yielded, or (2,1), but not both.

        Yields
        ------
        edge
            The edge as a pair of labels.

        """
        seen = set()

        for u, neighbors in self._adj.items():
            for v in neighbors:
                edge = frozenset((u, v))
                if edge not in seen:
                    seen.add(edge)
                    yield (u, v)
                else:
                    continue


class _DirectedEdgeView(_EdgeView):
    def __iter__(self):
        """Iterate through the edges.

        Yields
        ------
        edge
            The edge as an ordered pair of labels.

        """
        for u, neighbors in self._adj.items():
            for v in neighbors:
                yield (u, v)


class _Graph:
    """Base class for graph data structures."""

    def __init__(self, _edge_view_factory):
        self.adj = dict()
        self._number_of_edges = 0
        self._edge_view_factory = _edge_view_factory

    def __repr__(self):
        return f"<{self.__class__.__name__} with {len(self.nodes)} nodes and {len(self.edges)} edges>"

    def add_node(self, label):
        """
        Add a node with the given label.

        If the node already exists, nothing is done.

        Average case time complexity: Theta(1).

        Parameters
        ----------
        label
            The label of the node.

        """
        if label not in self.nodes:
            self.adj[label] = set()

    @property
    def nodes(self):
        """A view into the graph's nodes.

        Supports average case constant time node query. 

        Example
        -------
        >>> graph = UndirectedGraph()
        >>> graph.add_node('Red')
        >>> graph.add_node('Blue')
        >>> graph.add_node('Green')
        >>> 'Orange' in graph.nodes
        False
        >>> 'Red' in graph.nodes
        True
        >>> len(graph.nodes)
        3

        """
        return _NodesView(self.adj.keys())

    def arbitrary_node(self):
        """Return an arbitrary graph node. How the node is chosen is undefined.

        Takes Theta(1) time.

        Raises
        ------
        DoesNotExistError
            If the graph is empty.

        Example
        -------
        >>> graph = UndirectedGraph()
        >>> graph.add_node(1)
        >>> graph.add_node(2)
        >>> graph.add_node(3)
        >>> graph.arbitrary_node()
        2

        """
        try:
            return next(iter(self.nodes))
        except StopIteration:
            raise DoesNotExistError("The graph is empty.")

    @property
    def edges(self):
        """A view into the graph's edges.

        Supports average case constant time edge query.

        Example
        -------
        >>> graph = UndirectedGraph()
        >>> graph.add_edge('Red', 'Blue')
        >>> graph.add_edge('Blue', 'Green')
        >>> ('Red', 'Blue') in graph.edges
        True
        >>> ('Blue', 'Red') in graph.edges
        True
        >>> ('Red', 'Green') in graph.edges
        False
        
        """
        return self._edge_view_factory(self.adj, self._number_of_edges)


class UndirectedGraph(_Graph):
    def __init__(self, _edge_view_factory=_UndirectedEdgeView):
        super().__init__(_edge_view_factory)

    def add_edge(self, u_label, v_label):
        """Add an undirected edge to the graph.

        If the edge already exists, nothing is done.

        Average case time complexity: Theta(1).

        Parameters
        ----------
        u_label
            Label of one of the nodes in the edge.
        v_label
            Label of the other node in the edge.

        Notes
        -----
        If either of the nodes is not in the graph, the node is created.

        Raises
        ------
        ValueError
            If an attempt to add a self-loop is made. Undirected graphs do
            not have self-loops.

        """
        if u_label == v_label:
            raise ValueError("Undirected graphs have no self loops.")

        for x in {u_label, v_label}:
            if x not in self.adj:
                self.adj[x] = set()

        if (u_label, v_label) not in self.edges:
            self.adj[u_label].add(v_label)
            self.adj[v_label].add(u_label)
            self._number_of_edges += 1

    def remove_node(self, label):
        """Remove a node grom the graph.

        Average case time complexity: Theta(# of neighbors)
        
        Parameters
        ----------
        label
            The label of the node to be removed.

        Raises
        ------
        DoesNotExistError
            If the node is not in the graph.

        """
        try:
            neighbors = self.adj[label]
        except KeyError:
            raise DoesNotExistError(f'The node "{label}" does not exist.')

        for neighbor in self.adj[label]:
            self.adj[neighbor].discard(label)
            self._number_of_edges -= 1

        del self.adj[label]

    def remove_edge(self, u_label, v_label):
        """Remove the edge from the graph.

        Average case time complexity: Theta(1)

        Parameters
        ----------
        u_label
            The label of one of the nodes in the edge.
        v_label
            The label of the other node.

        Raises
        ------
        DoesNotExistError
            If the edge is not in the graph.

        """
        if (u_label, v_label) not in self.edges:
            raise DoesNotExistError(
                f'The edge "({u_label}, {v_label})" does not exist.'
            )

        self.adj[u_label].discard(v_label)
        self.adj[v_label].discard(u_label)
        self._number_of_edges -= 1

    def neighbors(self, label):
        """The neighbors of the node.

        Parameters
        ----------
        label
            The label of the node whose neighbors should be retrieved.

        Returns
        -------
        set
            The neighbors as a Python set. This set should not be modified.

        Note
        ----
        Since the return value is a set, there is no guarantee about the orders
        of the neighbors.

        """
        return self.adj[label]


class DirectedGraph(_Graph):
    def __init__(self, _edge_view_factory=_DirectedEdgeView):
        super().__init__(_edge_view_factory)
        self.back_adj = dict()

    def add_edge(self, u_label, v_label):
        """Add a directed edge to the graph.

        If the edge already exists, nothing is done.

        Average case time complexity: Theta(1).

        Parameters
        ----------
        u_label
            Label of the parent node.
        v_label
            Label of the child node.

        Note
        ----
        If either of the nodes is not in the graph, the node is created.

        """
        for x in {u_label, v_label}:

            if x not in self.adj:
                self.adj[x] = set()
            if x not in self.back_adj:
                self.back_adj[x] = set()

        self.adj[u_label].add(v_label)
        self.back_adj[v_label].add(u_label)
        self._number_of_edges += 1

    def remove_node(self, label):
        """Remove a node grom the graph.

        Average case time complexity: Theta(# of predecessors)
        
        Parameters
        ----------
        label
            The label of the node to be removed.

        Raises
        ------
        DoesNotExistError
            If the node is not in the graph.

        """
        if label not in self.nodes:
            raise DoesNotExistError(f'The node "{label}" does not exist.')

        # in case there is a self-loop, since we can't modify set while iterating
        if label in self.back_adj[label]:
            self.adj[label].discard(label)
            self.back_adj[label].discard(label)
            self._number_of_edges -= 1

        for parent in self.back_adj[label]:
            self.adj[parent].discard(label)
            self._number_of_edges -= 1

        self._number_of_edges -= len(self.adj[label])
        del self.adj[label]

    def remove_edge(self, u_label, v_label):
        """Remove the edge from the graph.

        Average case time complexity: Theta(1)

        Parameters
        ----------
        u_label
            The label of one of the nodes in the edge.
        v_label
            The label of the other node.

        Raises
        ------
        DoesNotExistError
            If the edge is not in the graph.

        """
        if (u_label, v_label) not in self.edges:
            raise DoesNotExistError(
                f'The edge "({u_label}, {v_label})" does not exist.'
            )

        self.adj[u_label].discard(v_label)
        self.back_adj[v_label].discard(u_label)
        self._number_of_edges -= 1

    def predecessors(self, label):
        """The predecessors of the node.

        Parameters
        ----------
        label
            The label of the node whose predecessors should be retrieved.

        Returns
        -------
        set
            The predecessors as a Python set. This set should not be modified.

        Note
        ----
        Since the return value is a set, there is no guarantee about the orders
        of the neighbors.


        """
        return self.back_adj[label]

    def successors(self, label):
        """The successors of the node.

        Parameters
        ----------
        label
            The label of the node whose successors should be retrieved.

        Returns
        -------
        set
            The successors as a Python set. This set should not be modified.

        Note
        ----
        Since the return value is a set, there is no guarantee about the orders
        of the neighbors.

        """
        return self.adj[label]

    def neighbors(self, label):
        """Alias of successors. Provided for convenience."""
        return self.successors(label)