Ticket #168: bp_matching.h

#ifndef SSP_HUNGARIAN_H
#define SSP_HUNGARIAN_H

#include <limits>
#include <list>
#include <algorithm>
#include <assert.h>
#include <queue>

#include <lemon/core.h>
#include <lemon/connectivity.h>
#include <lemon/bin_heap.h>

///\ingroup matching
///\file
///\brief Maximum weight matching algorithms in bipartite graphs.

namespace lemon {

  /// \ingroup matching
  ///
  /// \brief Maximum weight matching in (sparse) bipartite graphs
  ///
  /// This class implements a successive shortest path algorithm for finding
  /// a maximum weight matching in an undirected bipartite graph.
  /// Let \f$G = (X \cup Y, E)\f$ be an undirected bipartite graph. The 
  /// following linear program corresponds to a maximum weight matching
  /// in the graph \f$G\f$.
  ///
  /// \f$\begin{array}{rrcll} \
      \max & \displaystyle\sum_{(i,j) \in E} c_{ij} x_{ij}\\ \
      \mbox{s.t.} & \displaystyle\sum_{i \in X} x_{ij} & \leq & 1, \
      & \forall j \in \{ j^\prime \in Y \mid (i,j^\prime) \in E \}\\ \
      & \displaystyle\sum_{j \in Y} x_{ij} & \leq & 1, \
      & \forall i \in \{ i^\prime \in X \mid (i^\prime,j) \in E \}\\ \
      & x_{ij}        & \geq & 0, & \forall (i,j) \in E\\\end{array}\f$
  ///
  /// where \f$c_{ij}\f$ is the weight of edge \f$(i,j)\f$. The dual problem
  /// is:
  ///
  /// \f$\begin{array}{rrcll}\min & \displaystyle\sum_{v \in X \cup Y} p_v\\ \
      \mbox{s.t.} & p_i + p_j & \geq & c_{ij}, & \forall (i,j) \in E\\ \
      & p_v & \geq & 0, & \forall v \in X \cup Y \end{array}\f$
  ///
  /// A maximum weight matching is constructed by iteratively considering the
  /// vertices in \f$X = \{x_1, \ldots, x_n\}\f$. In every iteration \f$k\f$ 
  /// we establish primal and dual complementary slackness for the subgraph 
  /// \f$G[X_k \cup Y]\f$ where \f$X_k = \{x_1, \ldots, x_k\}\f$.
  /// So after the final iteration the primal and dual solution will be equal,
  /// and we will thus have a maximum weight matching. The time complexity of 
  /// this method is \f$O(n(n + m)\log n)\f$. 
  ///
  /// In case the bipartite graph is dense, it is better to use 
  /// \ref MaxWeightedDenseBipartiteMatching, which has a time complexity of
  /// \f$O(n^3)\f$.
  ///
  /// \tparam GR The undirected graph type the algorithm runs on.
  /// \tparam WM The type edge weight map. The default type is
  /// \ref concepts::Graph::EdgeMap "GR::EdgeMap<int>".
#ifdef DOXYGEN
  template <typename GR, typename WM>
#else
  template <typename GR, 
            typename WM = typename GR::template EdgeMap<int> >
#endif
  class MaxWeightedBipartiteMatching 
  {
    // TODO: 
    // - deal with orientation (edges need to be oriented from Y to X)
    //   (will be resolved when using a specific BpGraph class)
    // - swap X and Y in case |Y| < |X|, results in improved running times

  public:
    /// The graph type of the algorithm
    typedef GR Graph;
    /// The type of the edge weight map
    typedef WM WeightMap;
    /// The value type of the edge weights
    typedef typename WeightMap::Value Value;
    /// The type of the matching map
    typedef typename Graph::template NodeMap<typename Graph::Edge> MatchingMap;
    /// The type of a list of nodes
    typedef std::list<typename Graph::Node> NodeList;
  
  private:
    TEMPLATE_GRAPH_TYPEDEFS(Graph);

    typedef typename Graph::template NodeMap<bool> BoolMap;
    typedef typename Graph::template NodeMap<Node> AncestorMap;
    typedef typename Graph::template NodeMap<Value> PotMap;
    typedef typename Graph::template NodeMap<Value> DistMap;

    typedef typename Graph::template EdgeMap<bool> EdgeBoolMap;
    
    typedef Orienter<const Graph> OrientedGraph;
    typedef typename OrientedGraph::Arc OrientedArc;
    typedef typename OrientedGraph::ArcIt OrientedArcIt;
    typedef typename OrientedGraph::OutArcIt OrientedOutArcIt;
    typedef typename OrientedGraph::InArcIt OrientedInArcIt;
    typedef typename Graph::template NodeMap<OrientedArc> PredMap;

    typedef typename Graph::template NodeMap<int> HeapCrossRef;
    typedef BinHeap<Value, HeapCrossRef> Heap;

    const Graph& _graph;
    OrientedGraph* _pOrientedGraph;

    const WeightMap& _weight;
    BoolMap _bpMap;
    PotMap _pot;

    EdgeBoolMap* _pMatching;
    Value _matchingWeight;

    MatchingMap* _pMatchingMap;
    int _matchingSize;

    NodeList _X;
    NodeList _Y;

    bool isFree(const Node& v)
    {
      assert(_pOrientedGraph);
      // - a node x in X is free iff its in-degree is 0
      // - a node y in Y is free iff its out-degree is 0
      return (_bpMap[v] && OrientedInArcIt(*_pOrientedGraph, v) == INVALID) || 
        (!_bpMap[v] && OrientedOutArcIt(*_pOrientedGraph, v) == INVALID);
    }

    bool isFreeX(const Node& x)
    {
      assert(_bpMap[x] && _pOrientedGraph);
      // - a node x in X is free iff its in-degree is 0
      return OrientedInArcIt(*_pOrientedGraph, x) == INVALID;
    }

    bool isFreeY(const Node& y)
    {
      assert(!_bpMap[y] && _pOrientedGraph);
      // - a node y in Y is free iff its out-degree is 0
      return OrientedOutArcIt(*_pOrientedGraph, y) == INVALID;
    }

    void augmentPath(const Node& y, const PredMap& pred)
    {
      // M' = M ^ EP
      OrientedArc a = pred[y];
      while (a != INVALID)
      {
        _pOrientedGraph->reverseArc(a);

        if ((*_pMatching)[a])
        {
          _pMatchingMap->set(_pOrientedGraph->source(a), a);
          _pMatchingMap->set(_pOrientedGraph->target(a), a);
        }

        // arc has been reversed, so we need the target now
        a = pred[_pOrientedGraph->target(a)];
      }
    }

    void augment(const Node& x, DistMap& dist, PredMap& pred)
    {
      assert(isFreeX(x));

      /**
       * In case maxCardinality == false, we also need to consider 
       * augmenting paths starting from x and ending in a matched 
       * node x' in X. Augmenting such a path does *not* increase 
       * the cardinality of the matching. It may, however, increase 
       * the weight of the matching.
       *
       * Along with a shortest path starting from x and ending in 
       * a free vertex y in Y, we also determine x' such that
       *   y' = pred[x'],
       *   (pot[x] + pot[y'] - dist[x, y']) - w(y', x') is maximal
       *
       * Since (y', x') is part of the matching, 
       * by primal complementary slackness we have that
       *   pot[y'] + pot[x'] = w(y', x').
       *
       * Hence
       * x' = arg max_{x' \in X} { pot[x] + pot[y'] - dist[x, y']) -w(y', x') }
       *    = arg max_{x' \in X} { pot[x] - dist[x, y'] - pot[x'] }
       *    = arg max_{x' \in X} { -dist[x, y'] - pot[x'] }
       *    = arg min_{x' \in X} { dist[x, y'] + [x'] }
       *
       * We only augment x ->* x' if dist(x,y) > dist[x, y'] + pot[x']
       * Otherwise we augment x ->* y.
       */
    
      Value UB = _pot[x];
      NodeList visitedX, visitedY;
      
      // heap only contains nodes in Y
      HeapCrossRef heapCrossRef(_graph, Heap::PRE_HEAP);
      Heap heap(heapCrossRef);

      // add nodes adjacent to x to heap, and update UB
      visitedX.push_back(x);
      for (OrientedOutArcIt a(*_pOrientedGraph, x); a != INVALID; ++a)
      {
        const Node& y = _pOrientedGraph->target(a);
        Value dist_y = dist[x] + _pot[x] + _pot[y] - _weight[a];
        
        if (dist_y >= UB) 
          continue;

        if (isFreeY(y))
          UB = dist_y;

        dist[y] = dist_y;
        pred[y] = a;

        assert(heap.state(y) == Heap::PRE_HEAP);
        heap.push(y, dist_y);
      }

      Node x_min = x;
      Value minDist = 0, x_min_dist = _pot[x];

      while (true)
      {
        assert(heap.empty() || heap.prio() == dist[heap.top()]);

        if (heap.empty() || heap.prio() >= x_min_dist)
        {
          minDist = x_min_dist;

          if (x_min != x)
          {
            // we have an augmenting path between x and x_min 
            // that doesn't increase the matching size
            _matchingWeight += _pot[x] - x_min_dist;

            // x_min becomes free, and will always remain free
            (*_pMatchingMap)[x_min] = INVALID;
            augmentPath(x_min, pred);
          }
          break;
        }

        const Node y = heap.top();
        const Value dist_y = heap.prio();
        heap.pop();

        visitedY.push_back(y);
        if (isFreeY(y))
        {
          // we have an augmenting path between x and y
          augmentPath(y, pred);
          _matchingSize++;

          assert(_pot[y] == 0);
          _matchingWeight += _pot[x] - dist_y;
          
          minDist = dist_y;
          break;
        }
        else
        {
          // y is not free, so there *must* be only one arc pointing toward X
          OrientedArc a = OrientedOutArcIt(*_pOrientedGraph, y);
          assert(a != INVALID && ++OrientedOutArcIt(*_pOrientedGraph, y) == INVALID);

          const Node& x2 = _pOrientedGraph->target(a);
          pred[x2] = a;
          visitedX.push_back(x2);
          dist[x2] = dist_y; // matched edges have a reduced weight of 0

          // it must hold that x2 has only one incoming arc
          assert(OrientedInArcIt(*_pOrientedGraph, x2) != INVALID && 
            ++OrientedInArcIt(*_pOrientedGraph, x2) == INVALID);

          if (dist_y + _pot[x2] < x_min_dist)
          {
            x_min = x2;
            x_min_dist = dist_y + _pot[x2];

            // we have a better criterion now
            if (UB > x_min_dist) 
              UB = x_min_dist;
          }

          for (OrientedOutArcIt a2(*_pOrientedGraph, x2); a2 != INVALID; ++a2)
          {
            const Node& y2 = _pOrientedGraph->target(a2);
            Value dist_y2 = dist[x2] + _pot[x2] + _pot[y2] - _weight[a2];

            if (dist_y2 >= UB) 
              continue;

            if (isFreeY(y2))
              UB = dist_y2;

            if (heap.state(y2) == Heap::PRE_HEAP)
            {
              dist[y2] = dist_y2;
              pred[y2] = a2;
              heap.push(y2, dist_y2);
            }
            else if (dist_y2 < dist[y2])
            {
              dist[y2] = dist_y2;
              pred[y2] = a2;
              heap.decrease(y2, dist_y2);
            }
          }
        }
      }

      // restore primal complementary slackness
      // (reduced edge weight of matching edges should be 0)
      for (typename NodeList::const_iterator itX = visitedX.begin(); 
        itX != visitedX.end(); itX++)
      {
        const Node& x = *itX;
        assert(minDist - dist[x] >= 0);
        _pot[x] -= minDist - dist[x];
        assert(_pot[x] >= 0);
      }

      for (typename NodeList::const_iterator itY = visitedY.begin();
        itY != visitedY.end(); itY++)
      {
        const Node& y = *itY;
        assert(minDist - dist[y] >= 0);
        _pot[y] += minDist - dist[y];
        assert(_pot[y] >= 0);
      }
    }

  public:
    /// \brief Constructor
    ///
    /// Constructor.
    ///
    /// \param graph is the input graph
    /// \param weight are the edge weights
    MaxWeightedBipartiteMatching(const Graph& graph, const WeightMap& weight)
      : _graph(graph)
      , _pOrientedGraph(NULL)
      , _weight(weight)
      , _bpMap(graph)
      , _pot(graph, 0)
      , _pMatching(NULL)
      , _matchingWeight(0)
      , _pMatchingMap(NULL)
      , _matchingSize(0)
      , _X()
      , _Y()
    {
      // TODO: a swap might be beneficial if |_X| > |_Y|
      if (bipartitePartitions(_graph, _bpMap))
      {
        for (NodeIt v(_graph); v != INVALID; ++v)
        {
          if (_bpMap[v])
            _X.push_back(v);
          else
            _Y.push_back(v);
        }
      }
    }

    /// \brief Constructor
    ///
    /// Constructor. Instead of computing a bipartite partition, X and Y
    /// are used.
    ///
    /// \param graph is the input graph
    /// \param X is the first color class of the given bipartite graph
    /// \param Y is the second color class of the given bipartite graph
    /// \param weight are the edge weights
    ///
    /// \pre All edges need to be oriented from Y to X.
    MaxWeightedBipartiteMatching(const Graph& graph, 
                                 const NodeList& X, const NodeList& Y,
                                 const WeightMap& weight)
      : _graph(graph)
      , _pOrientedGraph(NULL)
      , _weight(weight)
      , _bpMap(graph, false)
      , _pot(graph, 0)
      , _pMatching(NULL)
      , _matchingWeight(0)
      , _pMatchingMap(NULL)
      , _matchingSize(0)
      , _X(X)//(X.size() <= Y.size() ? X : Y)
      , _Y(Y)//(X.size() > Y.size() ? X : Y)
    {
      // set _bpMap
      for (typename NodeList::const_iterator itX = _X.begin();
        itX != _X.end(); itX++)
      {
        _bpMap[*itX] = true;
      }
    }

    ~MaxWeightedBipartiteMatching()
    {
      delete _pMatching;
      delete _pOrientedGraph;
      delete _pMatchingMap;
    }

    /// \brief Initialize the algorithm
    ///
    /// This function initializes the algorithm.
    ///
    /// \param greedy indicates whether a nonempty initial matching 
    /// should be used; this might be faster in some cases.
    void init(bool greedy = false)
    {
      // reset data structures
      delete _pMatching;
      delete _pOrientedGraph;
      delete _pMatchingMap;

      _pMatchingMap = new MatchingMap(_graph, INVALID);
      _pMatching = new EdgeBoolMap(_graph, false);
      _pOrientedGraph = new OrientedGraph(_graph, *_pMatching);
      _matchingWeight = 0;
      _matchingSize = 0;

      // _pot[x] is set to maximum incident edge weight
      for (typename NodeList::const_iterator itX = _X.begin();
        itX != _X.end(); itX++)
      {
        const Node& x = *itX;

        Value maxWeight = 0;
        Edge e_max = INVALID;
        for (IncEdgeIt e(_graph, *itX); e != INVALID; ++e)
        {
          // _pot[y] = 0 for all y \in Y
          _pot[_graph.oppositeNode(x, e)] = 0;
         
          if (_weight[e] > maxWeight)
          {
            maxWeight = _weight[e];
            e_max = e;
          }
        }

        if (e_max != INVALID)
        {
          _pot.set(x, maxWeight);
          const Node& y = _graph.oppositeNode(x, e_max);
          if (greedy && isFreeY(y))
          {
            _matchingWeight += maxWeight;
            _matchingSize++;
            _pMatching->set(e_max, true);
            _pMatchingMap->set(x, e_max);
            _pMatchingMap->set(y, e_max);
          }
        }
      }
    }

    /// \brief Start the algorithm
    ///
    /// This function starts the algorithm.
    ///
    /// \pre \ref init() must have been called before using this function.
    void start()
    {
      DistMap dist(_graph, 0);
      PredMap pred(_graph, INVALID);

      for (typename NodeList::const_iterator itX = _X.begin();
        itX != _X.end(); itX++)
      {
        const Node& x = *itX;
        
        if (isFreeX(x))
          augment(x, dist, pred);
      }
    }

    /// \brief Run the algorithm.
    ///
    /// This method runs the \c %MaxWeightedBipartiteMatching algorithm.
    ///
    /// \param greedy indicates whether a nonempty initial matching 
    /// should be used; this might be faster in some cases.
    ///
    /// \note mwbm.run() is just a shortcut of the following code.
    /// \code
    ///   mwbm.init();
    ///   mwbm.start();
    /// \endcode
    void run(bool greedy = false)
    {
      init();
      start(greedy);
    }

    /// \brief Checks whether the solution is optimal
    ///
    /// Checks using the dual solution whether the primal solution is optimal.
    ///
    /// \return \c true if the solution is optimal.
    bool checkOptimality() const
    {
      assert(_pMatchingMap);

      /*
       * Primal:
       *   max  \sum_{i,j} c_{ij} x_{ij}
       *   s.t. \sum_i x_{ij} <= 1
       *        \sum_j x_{ij} <= 1
       *               x_{ij} >= 0
       *
       * Dual:
       *   min  \sum_j p_j + \sum_i r_i
       *   s.t. p_j + r_i >= c_{ij}
       *              p_j >= 0
       *              r_i >= 0
       *
       * Solution is optimal iff:
       * - Primal complementary slackness:
       *   - x_{ij} = 1  =>  p_j + r_i = c_{ij}
       * - Dual complementary slackness:
       *   - p_j != 0    =>  \sum_i x_{ij} = 1
       *   - r_i != 0    =>  \sum_j x_{ij} = 1
       */

      // check whether primal solution is feasible
      for (NodeIt n(_graph); n != INVALID; ++n)
      {
        int count = 0;
        for (IncEdgeIt e(_graph, n); e != INVALID; ++e)
        {
          if ((*_pMatching)[e])
            count++;
        }
        
        if (count > 1)
          return false;
      }

      // check whether dual solution is feasible
      for (NodeIt n(_graph); n != INVALID; ++n)
      {
        if (_pot[n] < 0) return false;
      }  
      for (EdgeIt e(_graph); e != INVALID; ++e)
      {
        if (_pot[_graph.u(e)] + _pot[_graph.v(e)] < _weight[e])
          return false;
      }

      // check primal complementary slackness
      for (EdgeIt e(_graph); e != INVALID; ++e)
      {
        // x_{ij} = 1  =>  p_j + r_i = c_{ij}
        if ((*_pMatching)[e] &&
            _pot[_graph.u(e)] + _pot[_graph.v(e)] != _weight[e])
          return false;
      }

      // check dual complementary slackness
      for (NodeIt n(_graph); n != INVALID; ++n)
      {
        // p_j != 0    =>  \sum_i x_{ij} = 1
        // r_i != 0    =>  \sum_j x_{ij} = 1
        if (_pot[n] != 0)
        {
          int count = 0;
          for (IncEdgeIt e(_graph, n); e != INVALID; ++e)
          {
            if ((*_pMatching)[e])
              count++;
          }
          if (count != 1)
            return false;
        }
      }
      
      return true;
    }

    /// \brief Return a const reference to the matching map.
    ///
    /// This function returns a const reference to a node map that stores
    /// the matching edge incident to each node.
    ///
    /// \pre init() must have been called before using this function.
    const MatchingMap& matchingMap() const
    {
      assert(_pMatchingMap);
      return *_pMatchingMap;
    }

    /// \brief Return the weight of the matching.
    ///
    /// This function returns the weight of the found matching.
    ///
    /// \pre init() must have been called before using this function.
    Value matchingWeight() const
    {
      return _matchingWeight;
    }

    /// \brief Return the number of edges in the matching.
    ///
    /// This function returns the number of edges in the matching.
    int matchingSize() const
    {
      return _matchingSize;
    }

    /// \brief Return \c true if the given edge is in the matching.
    ///
    /// This function returns \c true if the given edge is in the found
    /// matching.
    ///
    /// \pre init() must have been been called before using this function.
    bool matching(const Edge& e) const
    {
      assert(_pMatching);
      return _pMatching[e] != INVALID;
    }

    /// \brief Return the matching edge incident to the given node.
    ///
    /// This function returns the matching edge incident to the
    /// given node in the found matching or \c INVALID if the node is
    /// not covered by the matching.
    ///
    /// \pre init() must have been been called before using this function.
    Edge matching(const Node& n) const
    {
      assert(_pMatchingMap);
      return (*_pMatchingMap)[n];
    }

    /// \brief Return the mate of the given node.
    ///
    /// This function returns the mate of the given node in the found
    /// matching or \c INVALID if the node is not covered by the matching.
    ///
    /// \pre init() must have been been called before using this function.
    Node mate(const Node& n) const
    {
      assert(_pMatchingMap);

      if ((*_pMatchingMap)[n] == INVALID)
        return INVALID;
      else
        return _graph.oppositeNode(n, (*_pMatchingMap)[n]);
    }
  };

  /// \ingroup matching
  ///
  /// \brief Maximum weight matching in (dense) bipartite graphs
  ///
  /// This class provides an implementation of the classical Hungarian 
  /// algorithm for finding a maximum weight matching in an undirected
  /// bipartite graph. This algorithm follows the primal-dual schema.
  /// The time complexity is \f$O(n^3)\f$. In case the bipartite graph is 
  /// sparse, it is better to use \ref MaxWeightedBipartiteMatching, which 
  /// has a time complexity of \f$O(n^2 \log n)\f$ for sparse graphs.
#ifdef DOXYGEN
  template <typename GR, typename WM>
#else
  template <typename GR, 
            typename WM = typename GR::template EdgeMap<int> >
#endif
  class MaxWeightedDenseBipartiteMatching 
  {
  public:
    /// The graph type of the algorithm
    typedef GR Graph;
    /// The type of the edge weight map
    typedef WM WeightMap;
    /// The value type of the edge weights
    typedef typename WeightMap::Value Value;
    /// The type of the matching map
    typedef typename Graph::template NodeMap<typename Graph::Edge> MatchingMap;
    /// The type of a list of nodes
    typedef std::list<typename Graph::Node> NodeList;

  private:   
    TEMPLATE_GRAPH_TYPEDEFS(Graph);

    typedef typename Graph::template NodeMap<int> IdMap;
    typedef typename Graph::template NodeMap<bool> BoolMap;

    typedef std::vector<Node> ReverseIdVector;
    typedef std::vector<int> MateVector;
    typedef std::vector<Value> WeightVector;
    typedef std::vector<bool> BoolVector;

    class BpEdgeT 
    {
    private:
      Value _weight;
      Edge _edge;

    public:
      BpEdgeT()
        : _weight(0)
        , _edge(INVALID)
      {
      }

      void setWeight(Value weight)
      {
        _weight = weight;
      }

      Value getWeight() const
      {
        return _weight;
      }

      void setEdge(const Edge& edge)
      {
        _edge = edge;
      }

      const Edge& getEdge() const
      {
        return _edge;
      }
    };

    typedef std::vector<std::vector<BpEdgeT> > AdjacencyMatrixType;

    const Graph& _graph;

    const WeightMap& _weight;
    IdMap _idMap;
    MatchingMap _matchingMap;
    
    AdjacencyMatrixType _adjacencyMatrix;
    
    ReverseIdVector _reverseIdMapX;
    ReverseIdVector _reverseIdMapY;
    
    WeightVector _labelMapX;
    WeightVector _labelMapY;
    
    MateVector _mateMapX;
    MateVector _mateMapY;
    
    int _nX;
    int _nY;
    
    int _matchingSize;
    Value _matchingWeight;

    static const Value _minValue;
    static const Value _maxValue;

    void initAdjacencyMatrix(const NodeList& nodesX, const NodeList& nodesY)
    {
      // adjacency matrix has dimensions |X| * |Y|, 
      // every entry in this matrix is initialized to (0, INVALID)
      _adjacencyMatrix = AdjacencyMatrixType(nodesX.size(), 
        std::vector<BpEdgeT>(nodesY.size(), BpEdgeT()));
      
      // fill in edge weights by iterating through incident edges of nodes in X
      for (typename NodeList::const_iterator itX = nodesX.begin(); 
        itX != nodesX.end(); itX++)
      {
        Node x = *itX;
        int idX = _idMap[x];

        for (IncEdgeIt e(_graph, x); e != INVALID; ++e)
        {
          Node y = _graph.oppositeNode(x, e);
          int idY = _idMap[y];

          Value w = _weight[e];

          BpEdgeT& item = _adjacencyMatrix[idX][idY];
          item.setEdge(e);
          item.setWeight(w);

          // label of a node x in X is initialized to maximum weight 
          // of edges incident to x
          if (w > _labelMapX[idX])
            _labelMapX[idX] = w;
        } // for e
      } // for itX
    }

    void updateSlacks(WeightVector& slack, int x)
    {
      Value lx = _labelMapX[x];
      for (int y = 0; y < _nY; y++)
      {
        // slack[y] = min_{x \in S} [l(x) + l(y) - w(x, y)]
        Value val = lx + _labelMapY[y] - _adjacencyMatrix[x][y].getWeight();
        if (slack[y] > val)
          slack[y] = val;
      }
    }

    void updateLabels(const BoolVector& setS, const BoolVector& setT, WeightVector& slack)
    {
      // recall that slack[y] = min_{x \in S} [l(x) + l(y) - w(x,y)]

      // delta = min_{y \not \in T} (slack[y])
      Value delta = _maxValue;
      for (int y = 0; y < _nY; y++)
      {
        if (!setT[y] && slack[y] < delta)
          delta = slack[y];
      }

      // update labels in X
      for (int x = 0; x < _nX; x++)
      {
        if (setS[x])
          _labelMapX[x] -= delta;
      }

      // update labels in Y
      for (int y = 0; y < _nY; y++)
      {
        if (setT[y])
          _labelMapY[y] += delta;
        else
        {
          // update slacks
          // remember that l(x) + l(y) hasn't changed for x \in S and y \in T
          // the only thing that has changed is l(x) + l(y) for x \in S and y \not \in T
          slack[y] -= delta;
        }
      }
    }

    void buildMatchingMap()
    {
      _matchingWeight = 0;
      _matchingSize = 0;
      
      for (int x = 0; x < _nX; x++)
      {
        assert(_mateMapX[x] != -1);
        
        int y = _mateMapX[x];
        
        const Edge& e = _adjacencyMatrix[x][y].getEdge();
        _matchingMap[_reverseIdMapX[x]] = _matchingMap[_reverseIdMapY[y]] = e;
        
        if (e != INVALID)
        {
          // only edges that where present in the original graph count as a matching
          _matchingSize++;
          _matchingWeight += _weight[e];
        }
      }
    }

  public:
    /// \brief Constructor
    ///
    /// Constructor.
    ///
    /// \param graph is the input graph
    /// \param weight are the edge weights
    MaxWeightedDenseBipartiteMatching(const Graph& graph, 
                                      const WeightMap& weight)
      : _graph(graph)
      , _weight(weight)
      , _idMap(graph, -1)
      , _matchingMap(graph, INVALID)
      , _adjacencyMatrix()
      , _reverseIdMapX()
      , _reverseIdMapY()
      , _labelMapX()
      , _labelMapY()
      , _mateMapX()
      , _mateMapY()
      , _nX(0)
      , _nY(0)
      , _matchingSize(0)
      , _matchingWeight(0)
    {
    }

    /// \brief Initialize the algorithm
    ///
    /// This function initializes the algorithm.
    ///
    /// \return \c false if the provided graph is not bipartite.
    bool init()
    {
      // Compute bipartite partitions
      BoolMap bpMap(_graph);
      
      if (!bipartitePartitions(_graph, bpMap))
        return false;

      NodeList nodesX, nodesY;
      for (NodeIt v(_graph); v != INVALID; ++v)
      {
        if (bpMap[v])
          nodesX.push_back(v);
        else
          nodesY.push_back(v);
      }

      init(nodesX, nodesY);
      return true;
    }


    /// \brief Initialize the algorithm
    ///
    /// This function initializes the algorithm.
    ///
    /// \param nodesX is the first color class of the given bipartite graph
    /// \param nodesY is the second color class of the given bipartite graph
    void init(const NodeList& nodesX, const NodeList& nodesY)
    {
      // we need that |X| < |Y|
      const NodeList& X = nodesX.size() < nodesY.size() ? nodesX : nodesY;
      const NodeList& Y = nodesX.size() < nodesY.size() ? nodesY : nodesX;

      _nX = static_cast<int>(X.size());
      _nY = static_cast<int>(Y.size());

      // init matching is empty
      _mateMapX = MateVector(_nX, -1);
      _mateMapY = MateVector(_nY, -1);

      _reverseIdMapX = ReverseIdVector(_nX);
      _reverseIdMapY = ReverseIdVector(_nY);

      // labels of nodes in X are initialized to -INF, 
      // these will be updated during initAdjacencyMatrix()
      _labelMapX = WeightVector(_nX, _minValue);
      
      // labels of nodes in Y are initialized to 0,
      // these won't be updated during initAdjacencyMatrix()
      _labelMapY = WeightVector(_nY, 0);

      int x = 0;
      for (typename NodeList::const_iterator itX = X.begin(); 
        itX != X.end(); itX++, x++)
      {
        _idMap[*itX] = x;
        _reverseIdMapX[x] = *itX;
      }

      int y = 0;
      for (typename NodeList::const_iterator itY = Y.begin();
        itY != Y.end(); itY++, y++)
      {
        _idMap[*itY] = y;
        _reverseIdMapY[y] = *itY;
      }

      initAdjacencyMatrix(X, Y);
    }

    /// \brief Run the algorithm.
    ///
    /// This method runs the \c %MaxWeightedDenseBipartiteMatching algorithm.
    ///
    /// \note mwdbm.run() is just a shortcut of the following code.
    /// \code
    /// if (mwdbm.init()
    ///   mwdbm.start();
    /// \endcode
    void run()
    {
      if (init());
        start();
    }

    /// \brief Start the algorithm
    ///
    /// This function starts the algorithm.
    ///
    /// \pre \ref init() must have been called before using this function.
    void start()
    {
      // maps y in Y to x in X by which it was discovered
      MateVector discoveredY(_nY, -1);

      int matchingSize = 0;
      
      // pick a root
      for (int r = 0; r < _nX; )
      {
        assert(_mateMapX[r] == -1);

        // clear slack map, i.e. set all slacks to +INF
        WeightVector slack(_nY, _maxValue);

        // initially T = {}
        BoolVector setT(_nY, false);

        // initially S = {r}
        BoolVector setS(_nX, false);
        setS[r] = true;
        
        std::queue<int> queue;
        queue.push(r);

        updateSlacks(slack, r);

        bool augmented = false;
        while (!queue.empty() && !augmented)
        {
          int x = queue.front();
          queue.pop();

          for (int y = 0; y < _nY; y++)
          {
            if (!setT[y] && 
              _labelMapX[x] + _labelMapY[y] == _adjacencyMatrix[x][y].getWeight())
            {
              // y was (first) discovered by x
              discoveredY[y] = x;

              if (_mateMapY[y] != -1) // y is matched, extend alternating tree
              {
                int z = _mateMapY[y];

                // add z to queue if not in S
                if (!setS[z])
                {
                  setS[z] = true;
                  queue.push(z);
                  updateSlacks(slack, z);
                }
                
                setT[y] = true;
              }
              else // y is free, we have an augmenting path between r and y
              {
                matchingSize++;

                int cx, ty, cy = y;
                do {
                  cx = discoveredY[cy];
                  ty = _mateMapX[cx];

                  _mateMapX[cx] = cy;
                  _mateMapY[cy] = cx;

                  cy = ty;
                } while (cx != r);
                
                // we found an augmenting path, start a new iteration of the first for loop
                augmented = true;
                break; // break for y
              }
            }
          } // y \not in T such that (r,y) in E_l
        } // queue

        if (!augmented)
          updateLabels(setS, setT, slack);
        else
          r++;
      }

      buildMatchingMap();
    }

    /// \brief Return the weight of the matching.
    ///
    /// This function returns the weight of the found matching.
    ///
    /// \pre init() must have been called before using this function.
    Value matchingWeight() const
    {
      return _matchingWeight;
    }

    /// \brief Return the number of edges in the matching.
    ///
    /// This function returns the number of edges in the matching.
    int matchingSize() const
    {
      return _matchingSize;
    }

    /// \brief Return \c true if the given edge is in the matching.
    ///
    /// This function returns \c true if the given edge is in the found
    /// matching.
    ///
    /// \pre init() must have been been called before using this function.
    bool matching(const Edge& e) const
    {
      return _matchingMap[_graph.u(e)] != INVALID;
    }

    /// \brief Return the matching edge incident to the given node.
    ///
    /// This function returns the matching edge incident to the
    /// given node in the found matching or \c INVALID if the node is
    /// not covered by the matching.
    ///
    /// \pre init() must have been been called before using this function.
    Edge matching(const Node& n) const
    {
      return _matchingMap[n];
    }

    /// \brief Return the mate of the given node.
    ///
    /// This function returns the mate of the given node in the found
    /// matching or \c INVALID if the node is not covered by the matching.
    ///
    /// \pre init() must have been been called before using this function.
    Node mate(const Node& n) const
    {
      if (_matchingMap[n] == INVALID)
        return INVALID;
      else
        return _graph.oppositeNode(n, _matchingMap[n]);
    }

    /// \brief Return a const reference to the matching map.
    ///
    /// This function returns a const reference to a node map that stores
    /// the matching edge incident to each node.
    ///
    /// \pre init() must have been called before using this function.
    const MatchingMap& matchingMap() const
    {
      return _matchingMap;
    }

    /// \brief Checks whether the solution is optimal
    ///
    /// Checks using the dual solution whether the primal solution is optimal.
    ///
    /// \return \c true if the solution is optimal.
    bool checkOptimality() const
    {
      // We have to check three things:
      // - whether we really have a matching
      // - whether all nodes in X have a mate
      // - whether the labeling is feasible
      //
      // If all three conditions are satisfied then we know that we 
      // have a maximum weight bipartite matching (Kuhn-Munkres theorem)

      // check whether all nodes in X have a mate 
      // and that the two maps do not conflict
      for (int x = 0; x < _nX; x++)
      {
        if (!(_mateMapX[x] != -1 && _mateMapY[_mateMapX[x]] == x))
          return false;
      }

      // every x should be linked to by exactly one y
      for (int x = 0; x < _nX; x++)
      {
        int n = 0;
        for (int y = 0; y < _nY; y++)
        {
          if (_mateMapY[y] == x) n++;
        }

        if (n != 1) return false;
      }

      // every y should be linked to by at most one x
      for (int y = 0; y < _nY; y++)
      {
        int n = 0;
        for (int x = 0; x < _nX; x++)
        {
          if (_mateMapX[x] == y) n++;
        }

        if (n > 1) return false;
      }

      // the labeling should be feasible
      for (int x = 0; x < _nX; x++)
      {
        for (int y = 0; y < _nY; y++)
        {
          if (!(_adjacencyMatrix[x][y].getWeight() <=_labelMapX[x] + _labelMapY[y]))
            return false;
        }
      }

      return true;
    }
  };

  template<typename GR, typename WM>
  const typename WM::Value MaxWeightedDenseBipartiteMatching<GR, WM>::_maxValue = 
    std::numeric_limits<typename WM::Value>::max();

  template<typename GR, typename WM> 
  const typename WM::Value MaxWeightedDenseBipartiteMatching<GR, WM>::_minValue = 
    std::numeric_limits<typename WM::Value>::min();
}

#endif //SSP_HUNGARIAN_H