ぺんぎんメモ

プログラミングのメモです。たまに私生活のことや鬱っぽいことを書きます。

TopCoder SRM 771 Div1 Medium - TwoMonthScheduling

競プロ TopCoder

TopCoder SRM 771 Div1 Medium - TwoMonthScheduling

問題自体が複雑なので、理解するのに半日かかった。
そして、そこから実際に解き終わるのにも半日かかった。
とても本番のコンテストで解けそうにない。

ただ、問題が解けたということは知識量的には問題がないので、これからも問題を解き続けて少しずつ理解と考察のスピードを速めていきたい。

問題が若干複雑。

ひとつの仕事を終えるのに2ヶ月かかる。1ヶ月目の仕事量はA_i、2ヶ月目の仕事量はB_iである。1ヶ月目の仕事と2ヶ月目の仕事の間に空白期間を入れてはならない。必ず連続している必要がある。
このような仕事がN個与えられる。1ヶ月にできる仕事量の上限はworkersである。N個の仕事を終えるために最低何ヶ月必要かを求めてください。

という問題。

この問題は入力の受け取り方も少し特殊で、A_iとB_iを求めるためのコードを自分で書く必要がある。つまり、暗号化文字列を復号する、みたいな処理を自分で書く。けっこう面倒。なぜこんなことをさせるんだろう？素直にA_iとB_iを渡してくれればいいのに。もしかすると、配列の要素数が大きすぎると問題が起こる言語があるのかな。

この問題のDPは比較的素直。
だけど、何も考えずに実装すると計算量がO(N3)になって間に合わない。だから工夫する必要がある。具体的には、範囲の和を求めるのに累積和を、範囲を更新するのに遅延伝搬セグメント木を使う。

#include <bits/stdc++.h>

#define rep(i, n) for (int i = 0; i < int(n); i++)
#define rrep(i, n) for (int i = int(n) - 1; i >= 0; i--)
#define repc(i, n) for (int i = 0; i <= int(n); i++)
#define each(x, y) for (auto &x : y)

using namespace std;

using i32 = int;
using i64 = long long;
using vi32 = vector<i32>;
using vi64 = vector<i64>;

template <typename T, typename S>
struct LazySegmentTree {
  using F = function<T(T, T)>;
  using G = function<T(T, S, int)>;
  using H = function<S(S, S)>;

  vector<T> v;
  vector<S> z;
  F f;
  G g;
  H h;
  T e;
  S d;
  int n;

  LazySegmentTree(int size, F f, G g, H h, T e, S d) : f(f), g(g), h(h), e(e), d(d) {
    n = 1;
    while (n < size) n <<= 1;
    v.resize(n * 2, e);
    z.resize(n * 2, d);
  }
  
  void set(int k, T x) {
    v[k + n] = x;
  }

  void build() {
    rrep(i, n) v[i] = f(v[i * 2 + 0], v[i * 2 + 1]);
  }

  void propagate(int k, int l, int r) {
    v[k] = g(v[k], z[k], r - l);
    if (k < n) {
      z[k * 2 + 0] = h(z[k * 2 + 0], z[k]);
      z[k * 2 + 1] = h(z[k * 2 + 1], z[k]);
    }
    z[k] = d;
  }

  T update(int a, int b, S x, int k = 1, int l = 0, int r = -1) {
    if (r == -1) r = n;
    propagate(k, l, r);
    if (b <= l || r <= a) return v[k];
    if (a <= l && r <= b) { z[k] = x; propagate(k, l, r); return v[k]; }
    return v[k] = f(
      update(a, b, x, k * 2 + 0, l, (l + r) / 2),
      update(a, b, x, k * 2 + 1, (l + r) / 2, r)
    );
  }

  T query(int a, int b, int k = 1, int l = 0, int r = -1) {
    if (r == -1) r = n;
    propagate(k, l, r);
    if (b <= l || r <= a) return e;
    if (a <= l && r <= b) return v[k];
    return f(
      query(a, b, k * 2 + 0, l, (l + r) / 2),
      query(a, b, k * 2 + 1, (l + r) / 2, r)
    );
  }
};

struct TwoMonthScheduling {
  int minimumMonths(int workers, vi32 fm0, vi32 fm1, vi32 sm0, vi32 sm1) {
    int L0 = fm0.size(), L1 = fm1.size();
    vi32 fm(L0 * L1), sm(L0 * L1);
    rep(i1, L1) rep(i0, L0) {
      fm[i1 * L0 + i0] = min(workers, fm0[i0] ^ fm1[i1]);
      sm[i1 * L0 + i0] = min(workers, sm0[i0] ^ sm1[i1]);
    }
    int n = L0 * L1;

    vector<vi32> lst(2, vi32(n));
    rep(i, n) lst[0][i] = fm[i];
    rep(i, n) lst[1][i] = sm[i];

    vector<vi64> sums(2, vi64(n + 1));
    rep(i, 2) rep(j, n) sums[i][j + 1] += sums[i][j] + lst[i][j];

    vector<LazySegmentTree<int, int>> segs;
    repc(i, n) segs.push_back(LazySegmentTree<int, int>(
      n + 1,
      [](int a, int b) { return min(a, b); },
      [](int a, int b, int x) { return x >= 0 ? min(a, b) : -1e9; },
      [](int a, int b) { return min(a, b); },
      (int) 1e9, (int) 1e9)
    );

    segs[0].update(0, 1, 0);

    repc(i, n) repc(j, i) {
      i64 rest = (i64) workers - (sums[1][i] - sums[1][j]);
      if (rest < 0) continue;
      auto check = [&](int l, int r) {
        if (sums[0][r] - sums[0][l] > rest) return false;
        if (sums[1][r] - sums[1][l] > workers) return false;
        return true;
      };
      int ok = i, ng = n + 1;
      while (abs(ok - ng) > 1) {
        int md = (ok + ng) / 2;
        (check(i, md) ? ok : ng) = md;
      }
      segs[i].update(i, ok + 1, segs[j].query(i, i + 1) + 1);
    }

    return segs[n].query(n, n + 1);
  }
};

TopCoderでは、使われていない変数があると提出できないため、LazySegmentTreeのコンストラクタの第三引数の、無名関数の第三引数xを無理して使っている。

segs[n].query(n, n + 1)をsegs[n][n]と書けるようにしようか悩み中。[]は計算量O(1)感が凄いので実務とかでは避けたほうがいいと思うけど、競プロだから別にいいか、とか考えたりする。

そういえば、区間更新一点取得に特化した（定数倍高速化された）双対セグ木と呼ばれるものがあった。こちらを使うほうが今回の目的にマッチしている。遅延伝搬セグメント木は若干オーバースペック気味。

双対セグ木で[]演算子を定義するのはある程度自然だけど、遅延セグ木で[]演算子を定義していいかどうか悩む。まあいいや。定義しよう。

その他の感想

最初にO(N3)で実装し、実装し終えたあとで遅延セグ木を導入して計算量をO(N2 logN)に落とす、という順序で実装した。いきなりO(N2 logN)で実装すると、複数のバグが同時に発生して訳が分からなくなると思ったから。

この方針がよかったみたいで、遅延セグ木を久しぶりに使ったにも関わらず、コンパイルエラーを数回直すだけでO(N3)からO(N2 logN)への修正が完了した。出力がおかしくなることはなかった。

あと、自分ルールとして「ライブラリはすべて半開区間」というのがあって、それも実装時に助けになった。