階段関数の std::map による表現 (ACL Beginner Contest E - Replace Digits)

AtCoder の ACL Beginner Contest E - Replace Digits についての記事です．

問題へのリンク

atcoder.jp

経緯・考えたこと

コンテストのときには，この問題を解き始める時点でまだ80分以上残っていたのに，結局ACできませんでした．想定解は遅延伝搬セグメント木を使うというもので，それもチラッと頭には浮かんだのですが，モノイド演算が出てくると思えずに捨ててしまいました．出てきますよねえ...．

ともかくセグメント木は捨てて別の方向で考え始め，次のようにすれば良いのではないか，と思い至りました． 以下の記号を使います．

  • 1をk個並べた数 ( $ 1 + 10 + \cdots + 10^{k - 1} $ の値) を $C_k$ とする．
  • Sの位置は，右端 (1の位) を0，左端 ($10^{N-1}$の位) をN-1 とする．
  • Sの位置iの数を $x_i$ と書く．($ 1 \leq x_i \leq 9 $)
  • Tを，数が変化する位置と値のペアの集合とする．つまり，$x_{i} \neq x_{i - 1} $ なる $i$ に対する $(i, x_i)$ の全体を T とする．

T を持っていれば，「位置p から位置qまでdにせよ」というクエリに対して，Sの値が前回からどれだけ増減するかがわかります．$(i, x_i)$ と $(j, x_j)$ が T の中で隣り合っていて，$p \leq i < j \leq q$ だとすれば， この部分で $(d - x_i) \cdot C_{j - i} \cdot 10^i $ だけ増えるわけです．こういうのを全部足し合わせれば良い． Sの値を求めたら，Tを更新します．pとqの位置をTに加え，pからqの間にあるものはTから削除します．

計算量ですが，クエリで端点 p が導入されると，Tの要素が1増えますが，その後のクエリで，両端点の間に入らない限り，その要素は値の計算に必要ありません．両端点の間に入った場合には，次の瞬間にTから削除されます．つまり，この寄与は全体として$O(Q)$です．各クエリで，pとqの間に入るTの要素の境を探したり，pやqを挿入したり間の点を削除するのはそれぞれ $O(\log N)$ ですから，全体として $O(Q\log N)$ となり，十分間に合うはず (実際，間に合います)，ということで実装を始めました．

ところが，いつまでたってもコードが完成しません．C++のsetやmap の lower_bound や upper_bound を使ってコードを書くのですが，end() が返ってきたり，集合の中に等しい要素があったりなかったり，「○以上の最小の要素」はとれても，「○以下の最大の要素」は (すぐには) とれない，など，間違う要因が多くて，苦手です．一言で言ってしまえば，練習が足りないのですけれど．

階段関数の std::map による表現

ということで，コンテスト終了後になんとか動くコードを書いて得た教訓をメモしておきます．

Tはペアの集合ですから，map<int, int> で表現することにします．

(1) 番兵を置く．今回だと，T[-1] = 0, T[N] = 0. これらと合わせて，初期状態では，T は {(-1, 0), (0, 1), (N, 0)}.

(2) Sの値の変化を計算するときに，pからqまでいきなり足していくのは，私の技量では無理．次のようにする．

  • まず，$x_p$ と，$x_{q + 1}$ を求める．たとえば $x_p$ を求めるときには次のようにする．
    • p が T の key に入っていれば簡単．
    • そうでないときには，「p以下の最大の要素」を求めることになる． これは，reverse_iterator を使って，次のように実現できる:
auto it = T.lower_bound(p);
auto rit = reverse_iterator(it);
// *rit が求めるもの
  • T に，$(p, x_p)$ と $(q + 1, x_{q + 1})$ を追加する．これらの点はあっても害にならない．
  • p から q までの和を求める．pとq+1が T のキーになっているとわかっているから，場合分けなど無しで計算できる．

(3) 最後に，p と q+1 の間の要素を T から削除する．これも，p と q+1 は T のキーになっているとわかっているので，単純なループで実現できる:

auto goal = T.find(q+1);
auto it = T.find(p); it++;
while (it != goal) it = T.erase(it);

ACコード

(ライブラリコードが長いです) atcoder.jp