dlx.py

"""
Exact Cover Problem via Dancing Links
"""
# dlx.py
# Copyright (c) 2006,2008 Antti Ajanki <[email protected]>

# Heavily Modified Feb 2008, Tom Boothby
#  * Added a function which takes a Sage matrix and solves
#    the exact cover problem for it.
#  * Recursive search is now iterative
#  * Removed callback functionality
#  * Revamped the class to be a pythonic generator; new usage:
#        for cover in DLXMatrix(ones,initialsolution):
#            blah(cover)
#  * DOCUMENTATION AND TESTS GALORE HOLYCRAP 100% COVERAGE!!!!!1!!111@%^%*QQ!@E~

# DLXMatrix class is used to store and solve an exact cover problem
# with help of Dancing Links [1] technique by Donald Knuth. The
# problem can be stated as an attempt to leave out rows of a 0/1
# matrix until remaining matrix has exactly one 1 in each column.
# Knuth proposes a fast solution technique based on clever trick with
# double linked list.
#
# This implementation will return row numbers of the solution instead
# column indexes (like in the Knuth's paper).
#
# [1] Donald E Knuth, Dancing links, preprint, available at
# http://www-cs-faculty.stanford.edu/~knuth/preprints.html

# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.



ROOTNODE = 0

# Node's attributes
# LEFT, RIGHT, UP and DOWN are identifiers of corresponding neighbors
# INDEX is the (row) index of the node (None in the header nodes)
# COLUMN/COUNT is column index on a regular node and count of nodes on
# a header node
LEFT   = 0
RIGHT  = 1
UP     = 2
DOWN   = 3
COLUMN = 4
INDEX  = 5
COUNT  = 5


class DLXMatrix:
    def __init__(self, ones, initialsolution=None):
        """
        Solves the Exact Cover problem by using the Dancing Links algorithm
        described by Knuth.
        
        Consider a matrix M with entries of 0 and 1, and compute a subset
        of the rows of this matrix which sum to the vector of all 1's.
        
        The dancing links algorithm works particularly well for sparse
        matrices, so the input is a list of lists of the form: (note the
        1-index!)::

          [
           [1, [i_11,i_12,...,i_1r]]
           ...
           [m,[i_m1,i_m2,...,i_ms]]
          ]

        where M[j][i_jk] = 1.
        
        The first example below corresponds to the matrix::
        
           1110
           1010
           0100
           0001
        
        which is exactly covered by::
        
           1110
           0001
        
        and

        ::

           1010
           0100
           0001
        
        EXAMPLES::
        
            sage: from sage.combinat.dlx import *
            sage: ones = [[1,[1,2,3]]]
            sage: ones+= [[2,[1,3]]]
            sage: ones+= [[3,[2]]]
            sage: ones+= [[4,[4]]]
            sage: DLXM = DLXMatrix(ones,[4])
            sage: for C in DLXM:
            ...      print C
            [4, 1]
            [4, 2, 3]
        
        .. note::

           The 0 entry is reserved internally for headers in the
           sparse representation, so rows and columns begin their
           indexing with 1.  Apologies for any heartache this
           causes. Blame the original author, or fix it yourself.
        """
        if initialsolution is None: initialsolution = []
        self._cursolution = []
        self._nodes = [[0, 0, None, None, None, None]]
        self._constructmatrix(ones, initialsolution)
        self._level = 0
        self._stack = [(None,None)]

    def __eq__(self, other):
        r"""
        Return ``True`` if every attribute of
        ``other`` matches the attribute of
        ``self``.
        
        INPUT:
        
        
        -  ``other`` - a DLX matrix
        
        
        EXAMPLE::
        
            sage: from sage.combinat.dlx import *
            sage: M = DLXMatrix([[1,[1]]])
            sage: M == loads(dumps(M))
            True
        """
        if not isinstance(other, DLXMatrix):
            return False
        return self.__dict__ == other.__dict__

    def __iter__(self):
        """ 
        Returns self.
        
        TESTS::
        
            sage: from sage.combinat.dlx import *
            sage: M = DLXMatrix([[1,[1]]])
            sage: print M.__iter__() is M
            True
        """

        return self

    def _walknodes(self, firstnode, direction):
        """
        Generator for iterating over all nodes in given direction (not
        including firstnode).
        
        TESTS::
        
            sage: from sage.combinat.dlx import *
            sage: ones = [[1,[1,2,3]]]
            sage: ones+= [[2,[1,3]]]
            sage: ones+= [[3,[2]]]
            sage: ones+= [[4,[4]]]
            sage: DLX = DLXMatrix(ones,[4])
            sage: count = 0
            sage: for c in DLX._walknodes(ROOTNODE,RIGHT):
            ...       count += DLX._nodes[c][COUNT]
            ...       for d in DLX._walknodes(c,DOWN):
            ...           count -= 1
            sage: print count
            0
        """
        nodetable=self._nodes
        n = nodetable[firstnode][direction]
        while n != firstnode:
            yield n
            n = nodetable[n][direction]

    def _constructmatrix(self, ones, initialsolution=None):
        """
        Construct the (sparse) DLX matrix based on list 'ones'. The first
        component in the list elements is row index (which will be returned
        by solve) and the second component is list of column indexes of
        ones in given row.
        
        'initialsolution' is list of row indexes that are required to be
        part of the solution. They will be removed from the matrix.
        
        Note: rows and cols are 1-indexed - the zero index is reserved for
        the root node and column heads.
        
        TESTS::
        
            sage: from sage.combinat.dlx import *
            sage: ones = [[1,[1,2,3]]]
            sage: ones+= [[2,[1,3]]]
            sage: ones+= [[3,[2]]]
            sage: ones+= [[4,[4]]]
            sage: DLX = DLXMatrix([[1,[1]]])
            sage: DLX._constructmatrix(ones,[4])
            sage: c = DLX._nodes[ROOTNODE][RIGHT]
            sage: fullcount = 0
            sage: while c != ROOTNODE:
            ...       fullcount += DLX._nodes[c][COUNT]
            ...       d = DLX._nodes[c][DOWN]
            ...       while d != c:
            ...           bad = DLX._nodes[DLX._nodes[d][DOWN]][UP] != d
            ...           bad|= DLX._nodes[DLX._nodes[d][UP]][DOWN] != d
            ...           bad|= DLX._nodes[DLX._nodes[d][LEFT]][RIGHT] != d
            ...           bad|= DLX._nodes[DLX._nodes[d][RIGHT]][LEFT] != d
            ...           if bad:
            ...               raise RuntimeError, "Linked list inconsistent."
            ...           d = DLX._nodes[d][DOWN]
            ...       c = DLX._nodes[c][RIGHT]
            sage: print fullcount
            6
        """
        if initialsolution is None: initialsolution = []
        self._cursolution = []
        # LEFT, RIGHT, UP, DOWN, COLUMN, INDEX/COUNT
        self._nodes = [[ROOTNODE, ROOTNODE, None, None, None, None]]

        # optimization: local variables are faster
        nodetable = self._nodes
        ones.sort()
        pruneNodes = []
        headers = [ROOTNODE] # indexes of header nodes for faster access
        for r in ones:
            curRow = r[0]  # row index
            columns = r[1] # column indexes
            if len(columns) == 0: continue
            columns.sort()

            # Do we need more headers?
            while len(headers) <= columns[-1]:
                lastheader = headers[-1]
                newind = len(nodetable)
                nodetable.append([lastheader, ROOTNODE, newind, newind, None, 0])
                nodetable[ROOTNODE][LEFT] = newind
                nodetable[lastheader][RIGHT] = newind
                headers.append(newind)

            # Add new nodes to indexes newind..newind+len(columns)-1
            # LEFT and RIGHT links can be calculated when created the
            # node, only UP, DOWN and COUNT have to be updated when
            # adding new nodes
            newind = len(nodetable)
            for i, c in enumerate(columns):
                h = headers[c]
                l = newind + ((i-1) % len(columns))
                r = newind + ((i+1) % len(columns))
                nodetable.append([l, r, nodetable[h][UP], h, h, curRow])
                nodetable[nodetable[h][UP]][DOWN] = newind+i
                nodetable[h][UP] = newind+i
                nodetable[h][COUNT] += 1

            if curRow in initialsolution:
                pruneNodes.append(newind)


        # Remove columns that are required to be in the solution
        for n in pruneNodes:
            self._cursolution += [nodetable[n][INDEX]]
            self._covercolumn(nodetable[n][COLUMN])
            for j in self._walknodes(n, RIGHT):
                self._covercolumn(nodetable[j][COLUMN])

    def _covercolumn(self, c):
        """
        Performs the column covering operation, as described by Knuth's
        pseudocode::

           cover(c):
                i <- D[c]
                while i != c:
                    j <- R[i]
                    while j != i
                        D[U[j]] <- D[j]
                        U[D[j]] <- U[j]
                        N[C[j]] <- N[C[j]] - 1
                        j <- R[j]
                    i <- D[i]
        
        This is undone by the uncover operation.
        
        TESTS::
        
            sage: from sage.combinat.dlx import *
            sage: M = DLXMatrix([[1,[1,3]],[2,[1,2]],[3,[2]]])
            sage: one = M._nodes[ROOTNODE][RIGHT]
            sage: M._covercolumn(one)
            sage: two = M._nodes[ROOTNODE][RIGHT]
            sage: three = M._nodes[two][RIGHT]
            sage: print M._nodes[three][RIGHT] == ROOTNODE
            True
            sage: print M._nodes[two][COUNT]
            1
            sage: print M._nodes[three][COUNT]
            0
        """

        nodetable = self._nodes
        nodetable[nodetable[c][RIGHT]][LEFT] = nodetable[c][LEFT]
        nodetable[nodetable[c][LEFT]][RIGHT] = nodetable[c][RIGHT]
        for i in self._walknodes(c, DOWN):
            for j in self._walknodes(i, RIGHT):
                nodetable[nodetable[j][DOWN]][UP] = nodetable[j][UP]
                nodetable[nodetable[j][UP]][DOWN] = nodetable[j][DOWN]
                nodetable[nodetable[j][COLUMN]][COUNT] -= 1

    def _uncovercolumn(self, c):
        """
        Performs the column uncovering operation, as described by Knuth's
        pseudocode::

            uncover(c):
                i <- U[c]
                while i != c:
                    j <- L[i]
                    while j != i
                        U[j] <- U[D[j]]
                        D[j] <- D[U[j]]
                        N[C[j]] <- N[C[j]] + 1
                        j <- L[j]
                    i <- U[i]        
        
        This undoes by the cover operation since everything is done in the
        reverse order.
        
        TESTS::
        
            sage: from sage.combinat.dlx import *
            sage: M = DLXMatrix([[1,[1,3]],[2,[1,2]],[3,[2]]])
            sage: one = M._nodes[ROOTNODE][RIGHT]
            sage: M._covercolumn(one)
            sage: two = M._nodes[ROOTNODE][RIGHT]
            sage: M._uncovercolumn(one)
            sage: print M._nodes[two][LEFT] == one
            True
            sage: print M._nodes[two][COUNT]
            2
        """
        nodetable = self._nodes
        for i in self._walknodes(c, UP):
            for j in self._walknodes(i, LEFT):
                nodetable[nodetable[j][DOWN]][UP] = j
                nodetable[nodetable[j][UP]][DOWN] = j
                nodetable[nodetable[j][COLUMN]][COUNT] += 1
        nodetable[nodetable[c][RIGHT]][LEFT] = c
        nodetable[nodetable[c][LEFT]][RIGHT] = c

    def next(self):
        """
        Search for the first solution we can find, and return it.
        
        Knuth describes the Dancing Links algorithm recursively, though
        actually implementing it as a recursive algorithm is permissible
        only for highly restricted problems. (for example, the original
        author implemented this for Sudoku, and it works beautifully
        there)
        
        What follows is an iterative description of DLX::

            stack <- [(NULL)]
            level <- 0
            while level >= 0:
                cur <- stack[level]
                if cur = NULL:
                    if R[h] = h:
                        level <- level - 1
                        yield solution
                    else:
                        cover(best_column)
                        stack[level] = best_column
                else if D[cur] != C[cur]:
                    if cur != C[cur]:
                        delete solution[level]
                        for j in L[cur], L[L[cur]], ... , while j != cur:
                            uncover(C[j])
                    cur <- D[cur]
                    solution[level] <- cur
                    stack[level] <- cur
                    for j in R[cur], R[R[cur]], ... , while j != cur:
                        cover(C[j])
                    level <- level + 1
                    stack[level] <- (NULL)
                else:
                    if C[cur] != cur:
                        delete solution[level]
                        for j in L[cur], L[L[cur]], ... , while j != cur:
                            uncover(C[j])
                    uncover(cur)
                    level <- level - 1
        
        TESTS::
        
            sage: from sage.combinat.dlx import *
            sage: M = DLXMatrix([[1,[1,2]],[2,[2,3]],[3,[1,3]]])
            sage: while 1:
            ...     try:
            ...         C = M.next()
            ...     except StopIteration:
            ...         print "StopIteration"
            ...         break
            ...     print C
            StopIteration
            sage: M = DLXMatrix([[1,[1,2]],[2,[2,3]],[3,[3]]])
            sage: for C in M:
            ...       print C
            [1, 3]
            sage: M = DLXMatrix([[1,[1]],[2,[2,3]],[3,[2]],[4,[3]]])
            sage: for C in M:
            ...       print C
            [1, 2]
            [1, 3, 4]
        """
        nodetable = self._nodes # optimization: local variables are faster

        while self._level >= 0:
            c,r = self._stack[self._level]
            if c is None:
                if nodetable[ROOTNODE][RIGHT] == ROOTNODE:
                    self._level -= 1
                    return self._cursolution
                else:
                    c = nodetable[ROOTNODE][RIGHT]
                    maxcount = nodetable[nodetable[ROOTNODE][RIGHT]][COUNT]
                    for j in self._walknodes(ROOTNODE, RIGHT):
                        if nodetable[j][COUNT] < maxcount:
                            c = j
                            maxcount = nodetable[j][COUNT]
                    self._covercolumn(c)
                    self._stack[self._level] = (c,c)
            elif nodetable[r][DOWN] != c:
                if c != r:
                    self._cursolution = self._cursolution[:-1]
                    for j in self._walknodes(r, LEFT):
                        self._uncovercolumn(nodetable[j][COLUMN])
                r = nodetable[r][DOWN]
                self._cursolution += [nodetable[r][INDEX]]
                for j in self._walknodes(r, RIGHT):
                    self._covercolumn(nodetable[j][COLUMN])
                self._stack[self._level] = (c,r) 
                self._level += 1
                if len(self._stack) == self._level:
                    self._stack.append((None,None))
                else:
                    self._stack[self._level] = (None,None)
            else:
                if c != r:
                    self._cursolution = self._cursolution[:-1]
                    for j in self._walknodes(r, LEFT):
                        self._uncovercolumn(nodetable[j][COLUMN])
                self._uncovercolumn(c)
                self._level -= 1

        raise StopIteration



def AllExactCovers(M):
    """
    Utilizes A. Ajanki's DLXMatrix class to solve the exact cover
    problem on the matrix M (treated as a dense binary matrix).
    
    EXAMPLES::
    
        sage: M = Matrix([[1,1,0],[1,0,1],[0,1,1]])  #no exact covers
        sage: for cover in AllExactCovers(M):
        ...       print cover
        sage: M = Matrix([[1,1,0],[1,0,1],[0,0,1],[0,1,0]]) #two exact covers
        sage: for cover in AllExactCovers(M):
        ...       print cover
        [(1, 1, 0), (0, 0, 1)]
        [(1, 0, 1), (0, 1, 0)]
    """
    ones = []
    r = 1   #damn 1-indexing
    for R in M.rows():
        row = []
        for i in range(len(R)):
            if R[i]:
                row.append(i+1) #damn 1-indexing
        ones.append([r,row])
        r+=1
    for s in DLXMatrix(ones):
        yield [M.row(i-1) for i in s] #damn 1-indexing

def OneExactCover(M):
    """
    Utilizes A. Ajanki's DLXMatrix class to solve the exact cover
    problem on the matrix M (treated as a dense binary matrix).
    
    EXAMPLES::
    
        sage: M = Matrix([[1,1,0],[1,0,1],[0,1,1]])  #no exact covers
        sage: print OneExactCover(M)
        None
        sage: M = Matrix([[1,1,0],[1,0,1],[0,0,1],[0,1,0]]) #two exact covers
        sage: print OneExactCover(M)
        [(1, 1, 0), (0, 0, 1)]
    """

    for s in AllExactCovers(M):
        return s