test/library-checker/assignment.test.cpp

View this file on GitHub
Last update: 2022-06-25 20:33:37+09:00
Problem: https://judge.yosupo.jp/problem/assignment

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/assignment"

#include <cstdint>
#include <cstdio>

#include "src/graph/directed/flow/min_cost_flow.hpp"

int main() {
  using namespace workspace;

  size_t n;
  scanf("%zu", &n);
  min_cost_flow<int, int64_t> mcf(n * 2);
  for (size_t i = 0; i < n; i++) {
    mcf.supply(i);
    mcf.demand(i + n);
    for (size_t j = 0; j < n; j++) {
      int a;
      scanf("%d", &a);
      mcf.add_edge(i, j + n, 1, a);
    }
  }
  assert(mcf.run());
  printf("%lld\n", mcf.cost());
  for (size_t i = 0; i < n; i++) {
    for (auto &e : mcf[i]) {
      if (!e.capacity) printf("%zu ", e.head - n);
    }
  }
  puts("");
}

#line 1 "test/library-checker/assignment.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/assignment"

#include <cstdint>
#include <cstdio>

#line 2 "src/graph/directed/flow/min_cost_flow.hpp"

/**
 * @file min_cost_flow.hpp
 * @brief Minimum Cost Flow
 */

#include <algorithm>
#include <queue>

#line 2 "src/graph/directed/flow/base.hpp"

/**
 * @file base.hpp
 * @brief Flow Graph
 */

#include <cassert>
#include <numeric>
#include <tuple>
#include <vector>

namespace workspace {

template <class _Cap, class _Cost = void> class flow_graph {
 protected:
  class adjacency_impl;

 public:
  using container_type = std::vector<adjacency_impl>;
  using size_type = typename container_type::size_type;

  class unweighted_edge {
   public:
    size_type tail;  // Source
    size_type head;  // Destination
    _Cap capacity;   // Capacity
    _Cap flow;       // Flow

    unweighted_edge(size_type __s, size_type __d, const _Cap &__u = 1)
        : tail(__s), head(__d), capacity(__u), flow(0) {
      assert(!(capacity < static_cast<_Cap>(0))),
          assert(!(flow < static_cast<_Cap>(0)));
    }

    // tail, head, capacity, flow
    template <class _Os>
    friend _Os &operator<<(_Os &__os, const unweighted_edge &__e) {
      return __os << __e.tail << ' ' << __e.head << ' ' << __e.capacity << ' '
                  << __e.flow;
    }

   protected:
    unweighted_edge() = default;

    unweighted_edge(size_type __s, size_type __d, const _Cap &__u,
                    const _Cap &__f)
        : tail(__s), head(__d), capacity(__u), flow(__f) {}

    unweighted_edge make_rev() const { return {head, tail, flow, capacity}; }
  };

  class weighted_edge : public unweighted_edge {
   public:
    _Cost cost;  // _Cost

    weighted_edge(const unweighted_edge &__e, const _Cost &__c = 0)
        : unweighted_edge(__e), cost(__c) {}

    weighted_edge(size_type __s, size_type __d, const _Cap &__u = 1,
                  const _Cost &__c = 0)
        : unweighted_edge(__s, __d, __u), cost(__c) {}

    // tail, head, capacity, flow, cost
    template <class _Os>
    friend _Os &operator<<(_Os &__os, const weighted_edge &__e) {
      return __os << static_cast<unweighted_edge>(__e) << ' ' << __e.cost;
    }

   protected:
    weighted_edge() = default;

    weighted_edge make_rev() const {
      return {unweighted_edge::make_rev(), -cost};
    }
  };

  using edge = std::conditional_t<std::is_void<_Cost>::value, unweighted_edge,
                                  weighted_edge>;

 protected:
  struct edge_impl : edge {
    bool aux = false;
    edge_impl *rev = nullptr;

    edge_impl() = default;

    edge_impl(const edge &__e) : edge(__e) {}
    edge_impl(edge &&__e) : edge(__e) {}

    void push(_Cap __f) {
      edge::capacity -= __f;
      edge::flow += __f;
      if (rev) {
        rev->capacity += __f;
        rev->flow -= __f;
      }
    }

    edge_impl make_rev() {
      edge_impl __e = edge::make_rev();
      __e.aux = true;
      __e.rev = this;
      return __e;
    }
  };

 public:
  class adjacency {
   public:
    using value_type = edge;
    using reference = edge &;
    using const_reference = edge const &;
    using pointer = edge *;
    using const_pointer = const edge *;

    class iterator {
      edge_impl *__p;

     public:
      iterator(edge_impl *__p = nullptr) : __p(__p) {}

      bool operator!=(const iterator &__x) const { return __p != __x.__p; }

      bool operator==(const iterator &__x) const { return __p == __x.__p; }

      iterator &operator++() {
        do ++__p;
        while (__p->rev && __p->aux);
        return *this;
      }

      iterator operator++(int) {
        auto __cp = *this;
        do ++__p;
        while (__p->rev && __p->aux);
        return __cp;
      }

      iterator &operator--() {
        do --__p;
        while (__p->aux);
        return *this;
      }

      iterator operator--(int) {
        auto __cp = *this;
        do --__p;
        while (__p->aux);
        return __cp;
      }

      pointer operator->() const { return __p; }

      reference operator*() const { return *__p; }
    };

    class const_iterator {
      const edge_impl *__p;

     public:
      const_iterator(const edge_impl *__p = nullptr) : __p(__p) {}

      bool operator!=(const const_iterator &__x) const {
        return __p != __x.__p;
      }

      bool operator==(const const_iterator &__x) const {
        return __p == __x.__p;
      }

      const_iterator &operator++() {
        do ++__p;
        while (__p->rev && __p->aux);
        return *this;
      }

      const_iterator operator++(int) {
        auto __cp = *this;
        do ++__p;
        while (__p->rev && __p->aux);
        return __cp;
      }

      const_iterator &operator--() {
        do --__p;
        while (__p->aux);
        return *this;
      }

      const_iterator operator--(int) {
        auto __cp = *this;
        do --__p;
        while (__p->aux);
        return __cp;
      }

      const_pointer operator->() const { return __p; }

      const_reference operator*() const { return *__p; }
    };

    adjacency()
        : first(new edge_impl[2]), last(first + 1), __s(first), __t(first) {}

    ~adjacency() { delete[] first; }

    const_reference operator[](size_type __i) const {
      assert(__i < size());
      return *(first + __i);
    }

    size_type size() const { return __t - first; }

    auto begin() { return iterator{__s}; }
    auto begin() const { return const_iterator{__s}; }

    auto end() { return iterator{__t}; }
    auto end() const { return const_iterator{__t}; }

    /**
     * @brief Construct a new adjacency object
     *
     * @param __x Rvalue reference to another object
     */
    adjacency(adjacency &&__x) : first(nullptr) { operator=(std::move(__x)); }

    /**
     * @brief Assignment operator.
     *
     * @param __x Rvalue reference to another object
     * @return Reference to this object.
     */
    adjacency &operator=(adjacency &&__x) {
      delete[] first;
      first = __x.first, __x.first = nullptr;
      last = __x.last, __s = __x.__s, __t = __x.__t;
      return *this;
    }

   protected:
    edge_impl *first, *last, *__s, *__t;
  };

  using value_type = adjacency;
  using reference = adjacency &;
  using const_reference = adjacency const &;

 protected:
  class adjacency_impl : public adjacency {
   public:
    using _Base = adjacency;
    using _Base::__s;
    using _Base::__t;
    using _Base::first;
    using _Base::last;

    using iterator = edge_impl *;

    iterator push(const edge_impl &__e) {
      realloc();
      *__t = __e;
      if (__s->aux) ++__s;
      return __t++;
    }

    iterator push(edge_impl &&__e) {
      realloc();
      *__t = std::move(__e);
      if (__s->aux) ++__s;
      return __t++;
    }

    iterator begin() const { return first; }

    iterator end() const { return __t; }

    void realloc() {
      if (__t == last) {
        size_type __n(last - first);
        iterator loc = new edge_impl[__n << 1 | 1];
        __s += loc - first;
        __t = loc;
        for (iterator __p{first}; __p != last; ++__p, ++__t) {
          *__t = *__p;
          if (__p->rev) __p->rev->rev = __t;
        }
        delete[] first;
        first = loc;
        last = __t + __n;
      }
    }
  };

  // Only member variable.
  container_type graph;

 public:
  /**
   * @brief Construct a new flow graph object
   *
   * @param __n Number of vertices
   */
  flow_graph(size_type __n = 0) : graph(__n) {}

  /**
   * @brief Construct a new flow graph object
   *
   * @param __x Const reference to another object
   */
  flow_graph(const flow_graph &__x) : graph(__x.size()) {
    for (auto &&__adj : __x)
      for (auto &&__e : __adj) add_edge(__e);
  }

  /**
   * @brief Construct a new flow graph object
   *
   * @param __x Rvalue reference to another object
   */
  flow_graph(flow_graph &&__x) : graph(std::move(__x.graph)) {}

  /**
   * @brief Assignment operator.
   *
   * @param __x Const reference to another object
   * @return Reference to this object.
   */
  flow_graph &operator=(const flow_graph &__x) {
    return operator=(std::move(flow_graph{__x}));
  }

  /**
   * @brief Assignment operator.
   *
   * @param __x Rvalue reference to another object
   * @return Reference to this object.
   */
  flow_graph &operator=(flow_graph &&__x) {
    graph = std::move(__x.graph);
    return *this;
  }

  /**
   * @return Whether the graph is empty.
   */
  bool empty() const { return graph.empty(); }

  /**
   * @return Number of nodes.
   */
  size_type size() const { return graph.size(); }

  /**
   * @param node Node
   * @return Referece to the adjacency list of the node.
   */
  reference operator[](size_type node) {
    assert(node < size());
    return graph[node];
  }

  /**
   * @param node Node
   * @return Const referece to the adjacency list of the node.
   */
  const_reference operator[](size_type node) const {
    assert(node < size());
    return graph[node];
  }

  class iterator : public container_type::iterator {
    using _Base = typename container_type::iterator;

   public:
    using reference = adjacency &;
    using pointer = adjacency *;

    iterator(const _Base &__i) : _Base(__i) {}

    pointer operator->() const { return _Base::operator->(); }

    reference operator*() const { return _Base::operator*(); }
  };

  class const_iterator : public container_type::const_iterator {
    using _Base = typename container_type::const_iterator;

   public:
    using const_reference = const adjacency &;
    using const_pointer = const adjacency *;

    const_iterator(const _Base &__i) : _Base(__i) {}

    const_pointer operator->() const { return _Base::operator->(); }

    const_reference operator*() const { return _Base::operator*(); }
  };

  auto begin() { return iterator{graph.begin()}; }
  auto begin() const { return const_iterator{graph.begin()}; }

  auto end() { return iterator{graph.end()}; }
  auto end() const { return const_iterator{graph.end()}; }

  /**
   * @brief Add a node to the graph.
   *
   * @return Index of the node.
   */
  size_type add_node() { return add_nodes(1).front(); }

  /**
   * @brief Add some nodes to the graph.
   *
   * @param __n Number of nodes added
   * @return List of indices of the nodes.
   */
  std::vector<size_type> add_nodes(size_type __n) noexcept {
    std::vector<size_type> __nodes(__n);
    std::iota(__nodes.begin(), __nodes.end(), graph.size());
    graph.resize(graph.size() + __n);
    return __nodes;
  }

  /**
   * @brief Add a directed edge to the graph.
   *
   * @return Reference to the edge.
   */
  template <class... _Args>
  typename std::enable_if<std::is_constructible<edge, _Args...>::value,
                          edge &>::type
  add_edge(_Args &&...__args) {
    edge_impl __e = edge(std::forward<_Args>(__args)...);
    assert(__e.tail < size()), assert(__e.head < size());
    edge_impl *__p = graph[__e.tail].push(std::move(__e));
    // Careful with a self loop.
    if (__p->tail != __p->head)
      __p->rev = graph[__p->head].push(__p->make_rev());
    return *__p;
  }

  /**
   * @brief Add an undirected edge to the graph. Its cost must be non-negative.
   *
   * @return Reference to the edge.
   */
  template <class... _Args> edge &add_undirected_edge(_Args &&...__args) {
    edge_impl __e = edge(std::forward<_Args>(__args)...);
    assert(__e.tail < size()), assert(__e.head < size());
    __e.flow += __e.capacity;
    edge_impl *__p = graph[__e.tail].push(std::move(__e));
    // Careful with a self loop.
    if (__p->tail != __p->head) {
      edge_impl __r = __p->make_rev();
      __r.aux = false;
      __p->rev = graph[__p->head].push(std::move(__r));
    }
    return *__p;
  }

  template <class _Os>
  friend _Os &operator<<(_Os &__os, flow_graph const &__g) {
    for (const auto &__adj : __g)
      for (const auto &__e : __adj) __os << __e << "\n";
    return __os;
  }
};

}  // namespace workspace
#line 2 "lib/limits"

#include <limits>

namespace workspace {

template <class _Tp> struct numeric_limits : std::numeric_limits<_Tp> {};

#ifdef __SIZEOF_INT128__

template <> struct numeric_limits<__uint128_t> {
  constexpr static __uint128_t max() { return ~__uint128_t(0); }
  constexpr static __uint128_t min() { return 0; }
};

template <> struct numeric_limits<__int128_t> {
  constexpr static __int128_t max() {
    return numeric_limits<__uint128_t>::max() >> 1;
  }
  constexpr static __int128_t min() { return -max() - 1; }
};

#endif

}  // namespace workspace
#line 13 "src/graph/directed/flow/min_cost_flow.hpp"

namespace workspace {

/**
 * @brief Capacity Scaling Algorithm.
 *
 * @tparam _Cap Capacity type
 * @tparam _Cost Cost type
 */
template <class _Cap, class _Cost = _Cap>
class min_cost_flow : public flow_graph<_Cap, _Cost> {
  using _Base = flow_graph<_Cap, _Cost>;
  using typename _Base::edge_impl;

 public:
  using _Base::_Base;
  using _Base::add_edge;
  using typename _Base::edge;
  using typename _Base::size_type;

  /**
   * @brief Add a directed edge to the graph.
   *
   * @param __s Source
   * @param __d Destination
   * @param __l Lower bound of flow
   * @param __u Upper bound of flow
   * @param __c _Cost
   * @return Reference to the edge.
   */
  edge &add_edge(size_type __s, size_type __d, _Cap __l, _Cap __u, _Cost __c) {
    assert(!(__u < __l));
    b.resize(_Base::size());
    b[__s] -= __l, b[__d] += __l;
    auto &__e = _Base::add_edge(__s, __d, __u - __l, __c);
    __e.flow = __l;
    return __e;
  }

  /**
   * @brief Add an undirected edge to the graph.
   *
   * @return Reference to the edge.
   */
  template <class... _Args> edge &add_undirected_edge(_Args &&...__args) {
    auto &__e = static_cast<edge_impl &>(
        _Base::add_undirected_edge(std::forward<_Args>(__args)...));
    assert(!(__e.cost < 0));
    __e.rev->cost = __e.cost;
    return __e;
  }

  /**
   * @brief Increase the balance of a node.
   *
   * @param node
   * @param __f Default: 1
   */
  void supply(size_type node, _Cap __f = 1) {
    assert(node < _Base::size());
    b.resize(_Base::size());
    b[node] += __f;
  }

  /**
   * @brief Decrease the balance of a node.
   *
   * @param node
   * @param __f Default: 1
   */
  void demand(size_type node, _Cap __f = 1) {
    assert(node < _Base::size());
    b.resize(_Base::size());
    b[node] -= __f;
  }

  /**
   * @return Balance.
   */
  const auto &balance() const { return b; }

  /**
   * @param node Node
   * @return Balance of the node.
   */
  _Cap balance(size_type node) const { return b[node]; }

  /**
   * @return Potential. The dual solution.
   */
  const auto &potential() const { return p; }

  /**
   * @param node Node
   * @return Potential of the node.
   */
  _Cost potential(size_type node) const { return p[node]; }

  /**
   * @return _Cost of current flow.
   */
  _Cost cost() const { return current; }

  /**
   * @brief Run Capacity Scaling Algorithm.
   *
   * @return Whether a balanced flow exists.
   */
  bool run() {
    b.resize(_Base::size());
    p.resize(b.size());

    _Cap delta = 0;
    for (auto &&__adj : _Base::graph)
      for (auto &&__e : __adj) delta = std::max(delta, __e.capacity);

    if (delta == static_cast<_Cap>(0))
      if (std::any_of(b.begin(), b.end(),
                      [](_Cap __x) { return __x != static_cast<_Cap>(0); }))
        return false;

    parent.resize(b.size());

    while (static_cast<_Cap>(0) < delta) {
      delta /= 2;

      for (auto &&__adj : _Base::graph)
        for (auto &&__e : __adj)
          if (delta < __e.capacity && __e.cost < p[__e.head] - p[__e.tail]) {
            b[__e.tail] -= __e.capacity;
            b[__e.head] += __e.capacity;
            __e.push(__e.capacity);
          }

      sources.clear();
      sinks.clear();
      for (size_type __v = 0; __v != b.size(); ++__v)
        if (delta < b[__v])
          sources.emplace_back(__v);
        else if (b[__v] < -delta)
          sinks.emplace_back(__v);

      while (dual(delta)) {
        primal(delta);
        sources.erase(
            std::remove_if(sources.begin(), sources.end(),
                           [&](auto __v) { return !(delta < b[__v]); }),
            sources.end());
        sinks.erase(
            std::remove_if(sinks.begin(), sinks.end(),
                           [&](auto __v) { return !(b[__v] < -delta); }),
            sinks.end());
      }
    }

    current = 0;
    for (auto &&__adj : _Base::graph)
      for (auto &&__e : __adj)
        if (!__e.aux) current += __e.cost * __e.flow;

    return sources.empty() && sinks.empty();
  }

 protected:
  // _Cost of flow.
  _Cost current{};

  // Balance
  std::vector<_Cap> b;

  // The dual solution.
  std::vector<_Cost> p;

  std::vector<edge_impl *> parent;
  std::vector<size_type> sources, sinks;

  // Augment along the dual solution.
  void primal(_Cap delta) {
    for (auto __t : sinks)
      if (parent[__t]) {
        auto __f = -b[__t];
        auto __s = __t;
        while (parent[__s])
          __f = std::min(__f, parent[__s]->capacity), __s = parent[__s]->tail;
        if (delta < b[__s]) {
          __f = std::min(__f, b[__s]);
          b[__s] -= __f;
          b[__t] += __f;
          for (auto *__p = parent[__t]; __p; __p = parent[__p->tail]) {
            __p->push(__f);
            parent[__p->head] = nullptr;
          }
        }
      }
  }

  // Improve the dual solution.
  bool dual(_Cap delta) {
    std::fill(parent.begin(), parent.end(), nullptr);
    size_type reachable = 0;

    struct state {
      size_type __v;
      _Cost __d;
      state(size_type __v, _Cost __d) : __v(__v), __d(__d) {}
      bool operator<(const state &__x) const { return __x.__d < __d; }
    };

    std::priority_queue<state> __q;
    decltype(p) __nx(p.size(), numeric_limits<_Cost>::max());
    _Cost __ld = 0;

    for (auto __v : sources) {
      __nx[__v] = p[__v];
      __q.emplace(__v, 0);
    }

    while (!__q.empty()) {
      auto [__v, __d] = __q.top();
      __q.pop();
      if (__d + p[__v] != __nx[__v]) continue;
      __ld = __d;
      if (b[__v] < -delta && ++reachable == sinks.size()) break;
      for (auto &__e : _Base::graph[__v])
        if (delta < __e.capacity &&
            (__d = __nx[__v] + __e.cost) < __nx[__e.head]) {
          __q.emplace(__e.head, (__nx[__e.head] = __d) - p[__e.head]);
          parent[__e.head] = &__e;
        }
    }

    for (size_type __v = 0; __v != p.size(); ++__v)
      p[__v] = std::min(__nx[__v], p[__v] += __ld);

    return reachable;
  }
};

}  // namespace workspace
#line 7 "test/library-checker/assignment.test.cpp"

int main() {
  using namespace workspace;

  size_t n;
  scanf("%zu", &n);
  min_cost_flow<int, int64_t> mcf(n * 2);
  for (size_t i = 0; i < n; i++) {
    mcf.supply(i);
    mcf.demand(i + n);
    for (size_t j = 0; j < n; j++) {
      int a;
      scanf("%d", &a);
      mcf.add_edge(i, j + n, 1, a);
    }
  }
  assert(mcf.run());
  printf("%lld\n", mcf.cost());
  for (size_t i = 0; i < n; i++) {
    for (auto &e : mcf[i]) {
      if (!e.capacity) printf("%zu ", e.head - n);
    }
  }
  puts("");
}