第5回 ドワンゴからの挑戦状 本選 C - Interval and MST

問題ページ

解法

boruvka法を使って最小全域木問題を解く。各連結成分から他の連結成分へつなぐ最もコストが小さい辺をO(N)やO(NlogN)程度で求められるときに有効な方法になる。ある区間が他の区間と交差するのは4パターンに分類できる。

• その区間の右側と他の区間が交わる
  区間[l,r)に左端が含まれている場合これを満たす。このパターンの最もコストが小さい辺はr-(左端のmax)となる。左端が大きい上位2つの連結成分を保持しつつi=0から昇順に平面走査をすることで求められる。座圧しておくことで平面走査はO(N)で行える。内包する区間についてもこれは当てはまるが正しい値より大きい値となるため「その区間が他の区間を内包する」で求めるタイミングで正しい値に置き換わる。
  • iが右端になっている区間について交差する区間を考える
  • 今見ている区間と連結でなく左端が最も大きい区間を求める。左端が一番大きい区間が今見ている区間と連結であれば2番目に左端が大きい区間、そうでなければ最も左端が大きい区間となる。
  • iが左端になっている区間が存在した場合、左端が大きい上位2つの連結成分について更新する。
• その区間の左側と他の区間が交わる
  上記と同様に平面走査を行う。右端が小さい上位2つの連結成分を保持しつつi=(l[i],r[i]のうち最大)から降順に参照する。
• その区間が他の区間を内包する
  seg[i] = (左端がiの区間で区間長が短い上位2つ) としたセグ木を持つ。右端が小さい区間から順に見ていきセグ木の更新、区間[l,r)に内包される区間のうち連結成分が異なる区間長が短い区間を求める処理を順に行っていく。
• その区間が他の区間に外包される
  seg[i] = (右端がiの区間で区間長が短い上位2つ) としたセグ木を持つ。左端が小さい区間から順に見ていきセグ木の更新、区間[l,r)を外包する区間が存在すれば長さr-lで更新する処理を順に行っていく。

いずれのパターンもO(NlogN)で処理を行えるため全体でO(Nlog^2N)で求められる。

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
// #define int ll
using PII = pair<ll, ll>;

#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(x) x.begin(), x.end()

template<typename T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template<typename T> T &chmax(T &a, const T &b) { return a = max(a, b); }
template<typename T> bool IN(T a, T b, T x) { return a<=x&&x<b; }
template<typename T> T ceil(T a, T b) { return a/b + !!(a%b); }

template<typename T> vector<T> make_v(size_t a) { return vector<T>(a); }
template<typename T,typename... Ts>
auto make_v(size_t a,Ts... ts) {
  return vector<decltype(make_v<T>(ts...))>(a,make_v<T>(ts...));
}
template<typename T,typename V> typename enable_if<is_class<T>::value==0>::type
fill_v(T &t, const V &v) { t=v; }
template<typename T,typename V> typename enable_if<is_class<T>::value!=0>::type
fill_v(T &t, const V &v ) { for(auto &e:t) fill_v(e,v); }

template<class S,class T>
ostream &operator <<(ostream& out,const pair<S,T>& a){
  out<<'('<<a.first<<','<<a.second<<')'; return out;
}
template<typename T>
istream& operator >> (istream& is, vector<T>& vec){
  for(T& x: vec) {is >> x;} return is;
}
template<class T>
ostream &operator <<(ostream& out,const vector<T>& a){
  out<<'['; for(T i: a) {out<<i<<',';} out<<']'; return out;
}

int dx[] = {0, 1, 0, -1}, dy[] = {1, 0, -1, 0}; // DRUL
const int INF = 1<<30;
const ll LLINF = 1LL<<60;
const ll MOD = 1000000007;

struct UnionFind {
  vector<int> par, s;
  UnionFind(int n=2e5) { init(n); }
  void init(int n) { 
    s.assign(n, 1); par.resize(n); 
    iota(par.begin(), par.end(), 0);
  }
  int find(int x) {
    if(par[x] == x) return x;
    return par[x] = find(par[x]);
  }
  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if(x == y) return;
    if(s[x] < s[y]) par[x] = y, s[y] = s[x] + s[y];
    else par[y] = x, s[x] = s[x] + s[y];
  }
  bool same(int x, int y) { return find(x) == find(y); }
  int size(int x) { return s[find(x)]; }
};
 
template< typename T, typename F >
T boruvka(ll n, F f) {
  vector<ll> rev(n), belong(n);
  UnionFind uf(n);
  T ret = T();
  while(uf.size(0) != n) {
    ll ptr = 0;
    REP(i, n) {
      if(uf.find(i) == i) {
        belong[i] = ptr++;
        rev[belong[i]] = i;
      }
    }
    REP(i, n) belong[i] = belong[uf.find(i)];
    auto v = f(ptr, belong);
    bool update = false;
    REP(i, ptr) {
      if(~v[i].second && !uf.same(rev[i], rev[v[i].second])) {
        uf.unite(rev[i], rev[v[i].second]);
        ret += v[i].first;
        update = true;
      }
    }
    if(!update) return -1; // notice!!
  }
  return ret;
}

template <typename T>
struct segtree {
  int n;
  vector<T> dat;
  T d;
  using F = function<T(T,T)>;
  F f, g;

  segtree(int n_, F f_, F g_, T d_) : f(f_), g(g_), d(d_) {
    n = 1;
    while(n < n_) n *= 2;
    dat.assign(n*2, d); 
  }
  // [a, b)
  T query(int a, int b, int k, int l, int r) {
    if(r <= a || b <= l) return d;
    if(a <= l && r <= b) return dat[k];
    return f(query(a, b, 2*k+1, l, (l+r)/2),
             query(a, b, 2*k+2, (l+r)/2, r));
  }
  T query(int a, int b) {return query(a, b, 0, 0, n);}
  void update(int i, T v) {
    i += n-1;
    dat[i] = v;
    while(i > 0) {
      i = (i-1)/2;
      dat[i] = f(dat[i*2+1], dat[i*2+2]);
    }
  }
};
/**
* 点更新区間min d=INF, f=min(a,b), g=b\n
* 点更新区間max d=-INF, f=max(a,b), g=b\n
* 点加算区間和  d=0, f=a+b, g+=b
*/

signed main(void) 
{
  cin.tie(0);
  ios::sync_with_stdio(false);

  ll n;
  cin >> n;
  vector<ll> v, l(n), r(n);
  REP(i, n) {
    cin >> l[i] >> r[i];
    l[i]--, r[i]--;
    v.push_back(l[i]);
    v.push_back(r[i]);
  }
  sort(ALL(v));
  v.erase(unique(ALL(v)), v.end());
  const ll m = v.size();
  vector<vector<PII>> vl(m), vr(m);
  REP(i, n) {
    l[i] = lower_bound(ALL(v), l[i]) - v.begin();
    r[i] = lower_bound(ALL(v), r[i]) - v.begin();
    vl[l[i]].push_back({r[i], i});
    vr[r[i]].push_back({l[i], i});
  }

  // aが小さい方上位2つになるように更新
  function<pair<PII,PII>(pair<PII,PII>, PII)> 
  fmin = [](pair<PII,PII> a, PII b) {
    if(a.first.second == b.second) {
      a.first.first = min(a.first.first, b.first);
    } else if(a.second.second == b.second) {
      a.second.first = min(a.second.first, b.first);
      if(a.first > a.second) swap(a.first, a.second);
    } else {
      if(b < a.first) {
        a.second = a.first;
        a.first = b;
      } else if(b < a.second) {
        a.second = b;
      }
    }
    return a;
  };

  // aが大きい上位2つになるように更新
  function<pair<PII,PII>(pair<PII,PII>, PII)> 
  fmax = [](pair<PII,PII> a, PII b) {
    if(a.first.second == b.second) {
      a.first.first = max(a.first.first, b.first);
    } else if(a.second.second == b.second) {
      a.second.first = max(a.second.first, b.first);
      if(a.first < a.second) swap(a.first, a.second);
    } else {
      if(b > a.first) {
        a.second = a.first;
        a.first = b;
      } else if(b > a.second) {
        a.second = b;
      }
    }
    return a;
  };

  function<pair<PII,PII>(pair<PII,PII>, pair<PII,PII>)> 
  fmin2 = [&fmin](pair<PII,PII> a, pair<PII,PII> b) {
    a = fmin(a, b.first);
    a = fmin(a, b.second);
    return a;
  };

  function<vector<PII>(ll, vector<ll>)> 
  f = [&](ll sz, vector<ll> belong) {
    vector<PII> ret(sz, {LLINF, -1});

    // 右側で交差する区間について
    {
      // 左端が大きい方から上位2つを保持
      pair<PII,PII> top({-LLINF,-1}, {-LLINF,-1});
      REP(i, m) {
        for(auto p: vr[i]) {
          ll u = belong[p.second];
          // 今見ている区間の色と違う && 最も左端が大きい区間
          PII t = top.first.second == u ? top.second : top.first;
          // 交差していないか区間が存在していない
          if(t.first <= p.first) continue;
          if(t.second == -1) continue;
          chmin(ret[u], {v[i] - v[t.first], t.second});
        }
        for(auto p: vl[i]) top = fmax(top, {i, belong[p.second]});
      }
    }

    // 左側で交差する区間について
    {
      // 右端が小さい方から上位2つを保持
      pair<PII,PII> top({LLINF,-1}, {LLINF,-1});
      for(ll i=m-1; i>=0; --i) {
        for(auto p: vl[i]) {
          ll u = belong[p.second];
          // 今見ている区間の色と違う && 最も左端が大きい区間
          PII t = top.first.second == u ? top.second : top.first;
          // 交差していないか区間が存在していない
          if(t.first >= p.first) continue;
          if(t.second == -1) continue;
          chmin(ret[u], {v[t.first] - v[i], t.second});
        }
        for(auto p: vr[i]) top = fmin(top, {i, belong[p.second]});
      }
    }

    // 内包する区間について
    {
      segtree<pair<PII,PII>> seg(m, fmin2, fmin2, {{LLINF,-1},{LLINF,-1}});
      REP(i, m) {
        for(auto p: vr[i]) {
          // lの位置に区間の長さを更新
          seg.update(p.first, {{v[i] - v[p.first], belong[p.second]},{LLINF,-1}});
        }
        for(auto p: vr[i]) {
          pair<PII,PII> mi = seg.query(p.first, i);
          ll u = belong[p.second];
          PII t = mi.first.second == u ? mi.second : mi.first;
          chmin(ret[u], t);
        }
      }
    }

    // 外包される区間について
    {
      segtree<pair<PII,PII>> seg(m, fmin2, fmin2, {{LLINF,-1},{LLINF,-1}});
      REP(i, m) {
        for(auto p: vl[i]) {
          // rの位置に区間の長さを更新
          seg.update(p.first, {{v[p.first] - v[i], belong[p.second]},{LLINF,-1}});
        }
        for(auto p: vl[i]) {
          pair<PII,PII> mi = seg.query(p.first, m);
          ll u = belong[p.second];
          PII t = mi.first.second == u ? mi.second : mi.first;
          // 外包される区間があれば[i, p.first)の長さで更新
          if(t.second != -1) chmin(ret[u], {v[p.first] - v[i], t.second});
        }
      }
    }

    return ret;
  };

  cout << boruvka<ll, decltype(f)>(n, f) << endl;

  return 0;
}

実装がとても大変