graph/minimum-cost-arborescence.hpp

• View this file on GitHub
• Last update: 2024-08-10 13:03:16+09:00
• Include: #include "graph/minimum-cost-arborescence.hpp"

Depends on

• Skew Heap (data-structure/skew-heap.hpp)
• Union Find(Disjoint Set Union) (data-structure/union-find.hpp)
• グラフテンプレート (graph/graph-template.hpp)

Verified with

• verify/verify-yosupo-graph/yosupo-directed-mst.test.cpp

Code

#pragma once

#include "../data-structure/skew-heap.hpp"
#include "../data-structure/union-find.hpp"
#include "graph-template.hpp"

template <typename T>
Edges<T> MinimumCostArborescence(int N, int root, const Edges<T> &es) {
  using Heap = SkewHeap<T>;
  using Ptr = typename Heap::Ptr;
  UnionFind uf(N);
  vector<int> used(N, -1), from(N);
  vector<T> from_cost(N);
  vector<Ptr> come(N, nullptr);

  used[root] = root;
  vector<int> par_e(es.size(), -1), stem(N, -1), idxs;

  for (int i = 0; i < (int)es.size(); i++) {
    auto &e = es[i];
    come[e] = Heap::push(come[e], e.cost, i);
  }

  T costs = 0;

  for (int start = 0; start < N; start++) {
    if (used[start] != -1) continue;
    int cur = start;
    vector<int> chi_e;
    int cycle = 0;
    while (used[cur] == -1 || used[cur] == start) {
      used[cur] = start;
      if (come[cur] == nullptr) return {};
      int src = uf.find(es[come[cur]->idx].src);
      T cost = come[cur]->key + come[cur]->laz;
      int idx = come[cur]->idx;
      come[cur] = Heap::pop(come[cur]);
      if (src == cur) continue;

      from[cur] = src;
      from_cost[cur] = cost;
      if (stem[cur] == -1) stem[cur] = idx;
      costs += cost;
      idxs.push_back(idx);
      while (cycle) {
        par_e[chi_e.back()] = idx;
        chi_e.pop_back();
        --cycle;
      }
      chi_e.push_back(idx);

      if (used[src] == start) {
        int p = cur;
        do {
          if (come[p]) Heap::apply(come[p], -from_cost[p]);
          if (p != cur) {
            uf.unite(p, cur);
            Ptr newheap = Heap::meld(come[cur], come[p]);
            come[cur = uf.find(cur)] = newheap;
          }
          p = uf.find(from[p]);
          ++cycle;
        } while (p != cur);
      } else {
        cur = src;
      }
    }
  }

  vector<int> used_e(es.size());
  Edges<T> res;
  for (int _ = (int)idxs.size(); _--;) {
    int idx = idxs[_];
    if (used_e[idx]) continue;
    auto &e = es[idx];
    res.push_back(e);
    int x = stem[e.to];
    while (x != idx) {
      used_e[x] = true;
      x = par_e[x];
    }
  }
  return res;
}

#line 2 "graph/minimum-cost-arborescence.hpp"

#line 2 "data-structure/skew-heap.hpp"

template <typename T, bool isMin = true>
struct SkewHeap {
  struct Node {
    T key, laz;
    Node *l, *r;
    int idx;
    Node() = default;
    Node(const T &k, int i = -1)
        : key(k), laz(0), l(nullptr), r(nullptr), idx(i) {}
  };
  using Ptr = Node *;

  static void propagate(Ptr x) {
    if (x->laz == 0) return;
    if (x->l) x->l->laz += x->laz;
    if (x->r) x->r->laz += x->laz;
    x->key += x->laz;
    x->laz = 0;
  }

  static Ptr meld(Ptr x, Ptr y) {
    if (!x || !y) return x ? x : y;
    if (!comp(x, y)) swap(x, y);
    propagate(x);
    x->r = meld(x->r, y);
    swap(x->l, x->r);
    return x;
  }

  static Ptr alloc(const T &key, int idx = -1) { return new Node(key, idx); }

  static Ptr pop(Ptr x) {
    propagate(x);
    return meld(x->l, x->r);
  }

  static Ptr push(Ptr x, const T &key, int idx = -1) {
    return meld(x, alloc(key, idx));
  }

  static void apply(Ptr x, const T &laz) {
    x->laz += laz;
    propagate(x);
  }

 private:
  static inline bool comp(Ptr x, Ptr y) {
    if constexpr (isMin) {
      return x->key + x->laz < y->key + y->laz;
    } else {
      return x->key + x->laz > y->key + y->laz;
    }
  }
};

/**
 * @brief Skew Heap
 * @docs docs/data-structure/skew-heap.md
 */
#line 2 "data-structure/union-find.hpp"

struct UnionFind {
  vector<int> data;
  UnionFind(int N) : data(N, -1) {}

  int find(int k) { return data[k] < 0 ? k : data[k] = find(data[k]); }

  int unite(int x, int y) {
    if ((x = find(x)) == (y = find(y))) return false;
    if (data[x] > data[y]) swap(x, y);
    data[x] += data[y];
    data[y] = x;
    return true;
  }

  // f(x, y) : x に y をマージ
  template<typename F>
  int unite(int x, int y,const F &f) {
    if ((x = find(x)) == (y = find(y))) return false;
    if (data[x] > data[y]) swap(x, y);
    data[x] += data[y];
    data[y] = x;
    f(x, y);
    return true;
  }

  int size(int k) { return -data[find(k)]; }

  int same(int x, int y) { return find(x) == find(y); }
};

/**
 * @brief Union Find(Disjoint Set Union)
 * @docs docs/data-structure/union-find.md
 */
#line 2 "graph/graph-template.hpp"

template <typename T>
struct edge {
  int src, to;
  T cost;

  edge(int _to, T _cost) : src(-1), to(_to), cost(_cost) {}
  edge(int _src, int _to, T _cost) : src(_src), to(_to), cost(_cost) {}

  edge &operator=(const int &x) {
    to = x;
    return *this;
  }

  operator int() const { return to; }
};
template <typename T>
using Edges = vector<edge<T>>;
template <typename T>
using WeightedGraph = vector<Edges<T>>;
using UnweightedGraph = vector<vector<int>>;

// Input of (Unweighted) Graph
UnweightedGraph graph(int N, int M = -1, bool is_directed = false,
                      bool is_1origin = true) {
  UnweightedGraph g(N);
  if (M == -1) M = N - 1;
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    if (is_1origin) x--, y--;
    g[x].push_back(y);
    if (!is_directed) g[y].push_back(x);
  }
  return g;
}

// Input of Weighted Graph
template <typename T>
WeightedGraph<T> wgraph(int N, int M = -1, bool is_directed = false,
                        bool is_1origin = true) {
  WeightedGraph<T> g(N);
  if (M == -1) M = N - 1;
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    T c;
    cin >> c;
    if (is_1origin) x--, y--;
    g[x].emplace_back(x, y, c);
    if (!is_directed) g[y].emplace_back(y, x, c);
  }
  return g;
}

// Input of Edges
template <typename T>
Edges<T> esgraph([[maybe_unused]] int N, int M, int is_weighted = true,
                 bool is_1origin = true) {
  Edges<T> es;
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    T c;
    if (is_weighted)
      cin >> c;
    else
      c = 1;
    if (is_1origin) x--, y--;
    es.emplace_back(x, y, c);
  }
  return es;
}

// Input of Adjacency Matrix
template <typename T>
vector<vector<T>> adjgraph(int N, int M, T INF, int is_weighted = true,
                           bool is_directed = false, bool is_1origin = true) {
  vector<vector<T>> d(N, vector<T>(N, INF));
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    T c;
    if (is_weighted)
      cin >> c;
    else
      c = 1;
    if (is_1origin) x--, y--;
    d[x][y] = c;
    if (!is_directed) d[y][x] = c;
  }
  return d;
}

/**
 * @brief グラフテンプレート
 * @docs docs/graph/graph-template.md
 */
#line 6 "graph/minimum-cost-arborescence.hpp"

template <typename T>
Edges<T> MinimumCostArborescence(int N, int root, const Edges<T> &es) {
  using Heap = SkewHeap<T>;
  using Ptr = typename Heap::Ptr;
  UnionFind uf(N);
  vector<int> used(N, -1), from(N);
  vector<T> from_cost(N);
  vector<Ptr> come(N, nullptr);

  used[root] = root;
  vector<int> par_e(es.size(), -1), stem(N, -1), idxs;

  for (int i = 0; i < (int)es.size(); i++) {
    auto &e = es[i];
    come[e] = Heap::push(come[e], e.cost, i);
  }

  T costs = 0;

  for (int start = 0; start < N; start++) {
    if (used[start] != -1) continue;
    int cur = start;
    vector<int> chi_e;
    int cycle = 0;
    while (used[cur] == -1 || used[cur] == start) {
      used[cur] = start;
      if (come[cur] == nullptr) return {};
      int src = uf.find(es[come[cur]->idx].src);
      T cost = come[cur]->key + come[cur]->laz;
      int idx = come[cur]->idx;
      come[cur] = Heap::pop(come[cur]);
      if (src == cur) continue;

      from[cur] = src;
      from_cost[cur] = cost;
      if (stem[cur] == -1) stem[cur] = idx;
      costs += cost;
      idxs.push_back(idx);
      while (cycle) {
        par_e[chi_e.back()] = idx;
        chi_e.pop_back();
        --cycle;
      }
      chi_e.push_back(idx);

      if (used[src] == start) {
        int p = cur;
        do {
          if (come[p]) Heap::apply(come[p], -from_cost[p]);
          if (p != cur) {
            uf.unite(p, cur);
            Ptr newheap = Heap::meld(come[cur], come[p]);
            come[cur = uf.find(cur)] = newheap;
          }
          p = uf.find(from[p]);
          ++cycle;
        } while (p != cur);
      } else {
        cur = src;
      }
    }
  }

  vector<int> used_e(es.size());
  Edges<T> res;
  for (int _ = (int)idxs.size(); _--;) {
    int idx = idxs[_];
    if (used_e[idx]) continue;
    auto &e = es[idx];
    res.push_back(e);
    int x = stem[e.to];
    while (x != idx) {
      used_e[x] = true;
      x = par_e[x];
    }
  }
  return res;
}

Back to top page