Weighted balanced binary tree の平衡条件をWolfram Engineで

注: ガチ誰得記事

Weighted balanced binary treeという平衡二分探索木があります。詳細はwikipedia(Weight-balanced tree - Wikipedia)が詳しいのですが、ざっくり言うと

  • 「子のサイズが親のサイズの少なくとも $\alpha$ 倍」
  • もしくは同値な表現として「左右の子のサイズが高々 $\frac{1}{\alpha} - 1$ 倍しか違わない」

という平衡条件の木です。ここで、サイズというのは葉にのみ値を載せる木ならば葉の個数、全部の頂点に値を載せる木ならば(頂点の個数) + 1とします。

$\alpha$ を大きくすればするほど回転が増える代わりに強く平衡(=木の高さが低くなる)し、$\alpha$ を小さくすればするほど回転が減る代わりに平衡が弱くなります。

このサイズ情報はまず間違いなく競技プログラミングだと管理する(K番目のアクセスに必要なため)ので、赤黒木等と違い追加メモリが実質0と言う利点があります。カロリー0!

この木は強そうなのですが、うまく動く証明が難しいです。葉にのみ値を載せる木の場合のmerge関数はたとえば次のようになります。

Node* merge(Node* l, Node* r) {
    if (is_balanced(l->size, r->size)) return make_node(l, r);
    if (l->size > r->size) {
        auto x = l->l;
        auto y = merge(l->r, r);
        if (is_balanced(x->size, y->size)) return make_node(x, y);
        if (is_balanced(x->size, y->l->size) && is_balanced(x->size + y->l->size, y->r->size)) return make_node(make_node(x, y->l), y->r);
        return make_node(make_node(x, y->l->l), make_node(y->l->r, y->r));
    } else {
        ...
    }
}

このコードはどういうことかというと、lとrをマージしたい時、

  • lとrをそのままマージして問題ない: マージする
  • lがrに対してデカすぎてマージできない場合は、l->rとrを再帰的にmergeする。そしてl->lをx、merge(l->r, r)をyとする。
  • a = x, b = y->l->l, c = y->l->r, d = y->rとして、下図の3つの木のうち、少なくともどれか一つはうまくバランスしている。

という、大変魔法のようなマージ関数になります。

もちろん、任意の $\alpha$ についてこのmerge関数が動くわけではなく、wikipediaにあるように $0 \lt \alpha \lt 1 - \frac{1}{\sqrt{2}}$ ならば動きます。

今日はこの事実をWolfram Engine (Wolfram Engine) を利用して確かめていきましょう。

  • 変数 t を $\frac{1}{\alpha} - 1$ とします、つまり左右の子は高々 t 倍しかサイズが違わないという意味です。
  • 変数 a, b, c, d を上図の a, b, c, d のサイズとします
  • 変数 er のサイズとします。つまり l のサイズは a + b + c + d - e です。

まず、次のことが言えます

  • lr より t 倍以上大きい: (a+b+c+d-e) > t e
  • 元々 l->ll->rはバランスしていた: a <= t(b+c+d-e) && t a >= (b+c+d-e)
  • bcはバランスしている: b <= t c && t b >= c
  • b + cdはバランスしている: (b+c) <= t d && t(b+c) >= d

そしてこの条件下で

  • 図の一番上のマージが失敗する = (ab+c+dがバランスしていない): (t a < b+c+d || a > t(b+c+d))
  • 二番目のマージが失敗する = (ab+cがバランスしていない OR a+b+cdがバランスしていない): (t a < b+c || a > t(b+c) || t(a+b+c) < d || a+b+c > t d)
  • 三番目のマージが失敗する = (ab / cd / a+bc+d のどれかがバランスしていない): (t a < b || a > t b || t(a+b) < c+d || a+b > t(c+d) || t c < d || c > t d )

これを全てまとめてwolfram engineに投げると次のようになります

In[1]:= Resolve[Exists[{a, b, c, d, e}, 1 <= t && 0 <= a && 0 <= b && 0 <= c && 0 <= d && 0 <= e && (a+b+c+d-e) > t e && a <= t(b+c+d-e) && t a >= (b+c+d-e
) && b <= t c && t b >= c && (b+c) <= t d && t(b+c) >= d && (t a < b+c+d || a > t(b+c+d)) && (t a < b+c || a > t(b+c) || t(a+b+c) < d || a+b+c > t d) && (t
 a < b || a > t b || t(a+b) < c+d || a+b > t(c+d) || t c < d || c > t d )], t]

Out[1]= 1 <= t < 1 + Sqrt[2]

というわけで、$1 + \sqrt{2} \le t$、つまり$\alpha \le \frac{1}{2 + \sqrt{2}} = 1 - \frac{1}{\sqrt{2}}$ならば正しく動くことがわかりました。

結論: wolfram engineってすごい

データ構造をマージする一般的なテクを高速化するやや一般的なテク

知る人ぞ知る謎の知識みたいになってる気がしたのでメモ。

サイズ $A$ とサイズ $B$ のデータ構造 ($A < B$) が $O(A \log B)$ 時間でマージできる場合、合計 $O(N \log ^2 N)$ 時間で $N$ 個の要素がマージできます。競プロだと std::setstd::priority_queue などです。

ここで、もしもマージが $O(A (\log B - \log A + 1))$ に高速化できたならば、合計の計算量が $O(N \log N)$ になります。証明は普通のデータ構造をマージする一般的なテクとほぼ同じで、ある要素を含むデータ構造のサイズが $1 \to x \to y \to N$ と成長したなら、この要素にかかる計算量は $(\log x - \log 1 + 1) + (\log y - \log x + 1) + (\log N - \log z + 1) = O(\log N)$ 、といったノリです。

でも$O(A (\log B - \log A + 1))$ なんて謎の計算量なんて…と思うかもしれませんが、実はかなり多くのデータ構造がこの計算量を達成できます。すごいですね。

平衡二分木の merge / split が登場してややこしいので、データ構造をマージする操作を meld と呼ぶことにします

二分ヒープ

std::priority_queue の中身です。(自分で std::priority_queue 相当のものを実装する必要がありますが)実はpriority_queue + マージテクは $\log$ 一個です。でも $O(N \log ^2 N)$ のマージがそもそも爆速だからうまみが少なそう。

まず、二分ヒープを線形時間で構築するアルゴリズムを理解する必要があります。たとえば下のコードのようになります。

void down_heap(vector<int>& d, int u) {
    int n = int(d.size());
    while (2 * u + 1 < n) {
        int v = (2 * u + 1);
        if (v + 1 < n && h.d[v] < h.d[v + 1]) v++;
        if (h.d[u] >= h.d[v]) break;
        swap(h.d[u], h.d[v]);
        u = v;
    }
}

void build_heap(vector<int>& d) {
    int n = int(d.size());
    for (int i = n - 1; i >= 0; i--) {
        down_heap(d, i);
    }
    return h;
}

「ある頂点について、葉の方向に自分より大きい要素とswapしていく」down_heap関数を、葉から順に全ての頂点に呼ぶと二分ヒープが構築できます。計算量は各頂点の「葉までの距離」の総和になり、式をコネコネすると線形時間であることがわかります。

meld関数のコードは次のようになります。

void meld(vector<int>& h, vector<int>& other) {
    if (h.len() < other.len()) swap(h, other);
    if (other.empty()) return;

    int l = int(h.size()), r = l + int(other.size()) - 1;

    h.insert(h.end(), other.begin(), other.end());

    while (l) {
        l = (l - 1) / 2;
        r = (r - 1) / 2;
        for (int i = r; i >= l; i--) {
            down_heap(h, i);
        }
    }
}

これは何をしているのかというと、小さいほうのヒープの要素を適当に大きいヒープの末尾に追加した後、「新しく追加した要素を子孫に持つ要素」全てに対して先述のdown_heap関数を呼んでいるだけです。 計算量ですが、各頂点についてmin(other.size(), 葉までの距離)であること、そしてこれの総和が上記の $O(A (\log B - \log A + 1))$ になることが示せます。

std::set (by merge-split baseの平衡二分木)

std::setも、自作の平衡二分木で実装することで $O(A (\log B - \log A + 1))$が達成できます。まず、一般に merge-split base の平衡二分木ならなんでも可能な(ハズの)方法を紹介しますが、おそらく定数倍はとても悪いです。

$A \le \sqrt{B}$ の場合は愚直に $A$ 回 insertすればよいので、$\sqrt{B} < A$ としてよい。

  • サイズ $B$ の木を サイズ $B / A$ (ぐらい)の木に $A$ 分割する。
    • 愚直にsplitを使うと $O(A \log B)$ だが、例えば8分割したいなら (2分割 => それぞれを2分割 => それぞれを2分割)、のように、なるべくsplitする木のサイズが小さくなるように切っていくと $O(A (\log B - \log A + 1))$ になる
  • 分割した木それぞれに対して、サイズ $A$ の木の(対応する値の範囲の)要素を追加していく
    • 追加する要素が $B / A$ 個以下なら愚直にinsert
    • 追加する要素が $B / A$ 個以上なら(両方の木をvectorにする) => (std::merge) => (vectorを長さ $B / A$ のブロックに分割する) => (それぞれのブロックから、線形時間で新しい木を構築する)で、線形時間でマージする
  • 分割した木たちを(splitと同様に、いい感じに)マージして新しい木を作る。
    • サイズが $(B / A)$以上$2(B / A)$以下の木たちのマージになり、splitと同様の計算量解析が可能

で$O(A (\log B - \log A + 1))$ になっている…ハズ

std::set (by merge-split baseの多くの平衡二分木)

葉にだけ値を持たせる赤黒木など、サイズ $L, R (L < R)$の木のmerge (not meld)が $O(\log R - \log L + 1)$であることが保証可能な平衡二分木は割と簡単に実装できます。(参考: https://www2.ioi-jp.org/camp/2012/2012-sp-tasks/2012-sp-day4-copypaste-slides.pdf) このような平衡二分木なら、次のような割と素朴なmeld関数が上手くいくはずです。

Node* meld(Node* n, deque<int>& q) {
    while(q.size() && q.front() == n->val) q.pop_front(); // multisetにしたい場合はこの行を消す
    if (n->val < q.front()) return n;
    if (is_leaf(n)) {
        vector<int> v;
        while (q.size() && q.front() < n->val) {
            v.push_back(q.front()); q.pop_front();
        }
        return merge(build_tree(v), n); // build_tree(v): |v|時間で木を構築する関数
    }
    Node* l = meld(n->l, q);
    Node* r = meld(n->r, q);
    return merge(l, r); // (1)
}

Node* meld(Node* n, Node* m) {
    if (n->size < m->size) swap(n, m);
    deque<int> q = to_deque(m); // to_deque(m): mの要素を舐めて(sortedな)dequeを線形時間で生成する関数
    n = meld(n, q);
    return merge(n, build_tree(q));
}

まず、meld関数が呼ばれる回数は、n->sizeが $B / A$以上かどうかで場合分けすれば示せます。

  • n->size が $B / A$ 以上 : そもそもこういう頂点は $O(A)$ 個しかない
  • n->size が $B / A$ 未満: 高さが $O(\log (B / A)) = O(\log B - \log A + 1)$ なので合計 $O(A (\log B - \log A + 1))$

コード中の(1)の部分の計算量が本質です。ここで、merge関数の計算量により、meld前後で木のサイズが $k$ 倍に増えていたら (1) の部分の計算量が $O(\log k)$ であることを利用します。

  • 木のサイズが $2$ 倍以上に増えるような頂点数は?: 高々 $O(A)$
  • 木のサイズが $4$ 倍以上に増えるような頂点数は?: 高々 $O(A / 2)$
  • :

より、(*)の部分の計算量の合計は(meldが呼ばれた回数に加えて)高々 $O(A)$ です。

std::set (by splay-tree)

は愚直に小さいほうの集合を昇順(or 降順)にinsertしていくだけで $O(A (\log B - \log A + 1)$ って噂を聞いたんですが、本当かわかりません。

std::set (by treap)

は割と容易に $O(A (\log B - \log A + 1)$ のmeldが書けるって聞きました。いかがでしたか

UCUP 2-11 (Nanjing) E: Extending Distance / 最小費用流の双対

久しぶりですごい時間がかかったのでメモ

大体次の問題。

  • $N$ 頂点 $M$ 辺の有向グラフと非負整数 $K$ が与えられる。各辺には非負整数の重みが付いていて、辺 $e$ の重みを $d_e$ とする。$K$ は $1 \to N$ の最短距離より小さくない。 $1$ 円払うと好きな辺の重みを $1$ 増やせる時、頂点 $1$ から $N$ の最短距離を $K$ にするために必要な最小コストは? 構築あり: 各辺について何回伸ばすかも復元する必要あり

まず、この問題はほぼ最小費用循環流の双対問題そのものなので、構築がなければ難しくない。というか大体コレ J - Longest Shortest Path

最小費用循環流:

$N$ 頂点 $M$ 辺の有向グラフが与えられる。各辺 $e$ には非負の容量 $c_e$ と(非負とは限らない)コスト $d_e$ が付いている。辺ごとに $0 \le f_e \le c_e$を満たす流量 $f_e$ を割り当てる。頂点ごとに流量の出入りが等しい必要がある。 $\sum_e f_e d_e$を最小化せよ。

最小費用循環流の双対問題:

$N$ 頂点 $M$ 辺の有向グラフが与えられる。各辺 $e$ には非負整数 $c_e$ と整数 $d_e$ が付いている。頂点ごとにポテンシャル $p_v$ を割り当て、$- \sum_{e = (u \to v)} c_e \max(0, p_v - p_u - d_e)$ を最大化せよ。

この2つの問題はLP双対から得られる問題である。そのため、同じグラフに対するこの2つの問題の答えは必ず一致する。なお、(答えが $-1$ 倍されるが)後者の目的関数は「$\sum_{e = (u \to v)} c_e \max(0, p_v - p_u - d_e)$ を最小化」としたほうが見通しがいいと思う。参考: 双対性 | PPT

今回の問題

は、牛ゲー(最短経路問題の双対)に思いを馳せると、$(u, v, c, d) = (N, 1, \inf, -K), (1, N, \inf, K)$ の $2$ 辺を追加して、最小費用循環流の双対問題を解けばよいことがわかる。追加した $2$ 辺は $p_N - p_1 = K$という制約を追加することに対応する。これは普通の逐次最短路法+少しの工夫でもいいし、強力なアルゴリズムを用いてもいい (参考: https://judge.yosupo.jp/problem/min_cost_b_flow 周りの諸々)。つまり、この問題の本質は「最小費用循環流の双対問題の解 $p_v$ をどうやって復元するのか?」という点になる。

$p_v$ の復元方法

結論から言うと、大体のライブラリは内部にこの $p$ を変数として持っているし、最小費用循環流を流せるだけ流した後にベルマンフォード法を行えば直接求まる。どういうことだろうか?

まず、最小費用循環流の基礎として、次が成立する。

  • 最小費用循環流において、ある流量 $f$ が最適解である $\Leftrightarrow$ 残余グラフに負閉路が存在しない

($\Rightarrow$)は自明、($\Leftarrow$)は自明ではないけど、現在の流量 $f$ と最適解の流量 $f'$ の差分に注目すると示せる。

また、負閉路が存在しない、またその時のみ、残余グラフの(容量が $0$ でない)各辺について $p_v - p_u - d_e \le 0$ を満たすポテンシャル $p$ が存在する。これはベルマンフォード法などで計算可能。また、このポテンシャルは多くのライブラリが内部に直接変数として持っている。

逆に言うと、流量条件を満たす流量 $f$ と、$f$ の残余グラフにおいて $p_v - p_u - d_e \le 0$を満たすポテンシャル $p$ が見つけられたなら、そのペア $(f, p)$ は $f$ が最適解であるという証拠になる。…というわけで https://judge.yosupo.jp/problem/min_cost_b_flow では、流量 $f$ に加えてこのポテンシャル $p$ も復元させている。

2つのポテンシャル

この記事では $2$ つのポテンシャル $p$ が登場している。ひとつは「最小費用循環流の双対問題」で紹介した $p$ であり、これが今回の問題を解くために必要なものである。もう一つは最小費用循環流の最適解の残余グラフにベルマンフォード法を行うと計算可能な $p$ である。後者が計算可能なのはわかったが、どうやって前者を計算すればいいのか?実は後者の $p$ をそのまま前者の $p$ にすればいい(ええー!)

今までの話を再度まとめる。

計算可能なことがわかっているもの : 次の条件を満たすペア $(f, p)$

  • 各辺 $e$ について、$0 \le f_e \le c_e$
  • 各頂点について、流量の出入りの総和が等しい
  • $f_e > 0$ を満たす辺 $e = (u \to v)$ について、$p_u - p_v + d_e \le 0$
  • $f_e < c_e$ を満たす辺 $e = (u \to v)$ について、$p_v - p_u - d_e \le 0$

ここで、「最小費用循環流」の答えは $\sum_e f_e d_e$ となる。

計算したいもの: $\sum_{e = (u \to v)} c_e \max(0, q_v - q_u - d_e)$ が「最小費用循環流の答えの $-1$ 倍」となる $q$。

証明したいこと: $p$が求める $q$ のうちひとつであること。つまり、$\sum_{e = (u \to v)} c_e \max(0, p_v - p_u - d_e)$ = $- \sum_e f_e d_e$ を示せればよい。

証明

$f_e < c_e$ならば、$p_v - p_u - d_e \le 0$、つまり $\max(0, p_v - p_u - d_e) = 0$であるので、$\sum_{e = (u \to v)} c_e \max(0, p_v - p_u - d_e)$

$= \sum_{e = (u \to v)} f_e \max(0, p_v - p_u - d_e)$ となる。

また、$f_e > 0$ならば $p_v - p_u - d_e \ge 0$ なので、

$= \sum_{e = (u \to v)} f_e (p_v - p_u - d_e)$ となる。

そして、$f$ を単純閉路に分解して考えることで

$= \sum_{e = (u \to v)} - f_e d_e$

が言えるため、示せた。

結論

https://judge.yosupo.jp/problem/min_cost_b_flow を解くライブラリを準備(or 何らかの方法で入手)しておけば、張ると通る

こういう出題は許されるのか?⇒競プロで学ぶ統計学でした?

AtCoderではともかく、ほかのジャッジでは想定解の正当性の保証がされていない出題が行われることがあります。 想定解の正当性の保証が出来ていなくても、ジャッジが大量にケースを入れてちゃんと動いたからOK!という出題は可能なのか考えてみます。 例えば色々こだわると、次のような出題になると思います。

問題

長さ10000の数列が与えられる。各要素はそれぞれ1以上10000以下の整数である。次の条件を満たす部分列を探し、見つけたら出力せよ。

ジャッジ方法

あなたの提出は次のようにジャッジされる。

  • 以下を独立に100回行う。あなたのコードが少なくとも30ケースに対して正しい部分列を出力していた場合、ACとする。
    • ランダムケースを一様に$[1,10000]^{10000}$からジャッジが生成する。これは事前に作成されたものではなく、ジャッジごとに新たに作り直される。
    • このランダムケースをあなたのコードに入力し、部分列を出力した場合それが正当かを検証する

解説

解法コードについて、$[1, 10000]^{10000}$について正しい部分列を出力するケースの割合を $p$ とする。これはその解法コードのランダムケースに対して正しく動く確率である。

もし確率$p \geq 0.5$で正しい部分列を出力する解があれば、そのコードは<十分高い確率>でACする。

私たちジャッジは次のヒューリスティックを行う解答を用意しました。

そして、この解答に同様にランダムケースを5000ケース入力したら、4000ケースに対して正しい部分列を出力した。もしこの想定解法が正しく動く確率$p$が$p < 0.5$ならば、ランダムケースを5000ケース入れて4000ケースに対して正しく動く確率は<ハチャメチャ低い確率>である。よって、このジャッジ解は$p \geq 0.5$であると考えられる。

疑問

  • このような問題は許容されるのか?
  • 「よって、このジャッジ解は$p \geq 0.5$であると考えられる。」これはどういうことなの?
  • 許容される場合、ジャッジが事前に行うべきテストの量(今回でいうと4000/5000AC)は何らかの式で計算できるのか?

考察: 大嘘解法の可能性は消せない

例えば想定解が$p=0.01$であったとしても、確率$0.01^{5000}$で5000ケースに連続ACする。つまり、想定解が大嘘である可能性というのは0には出来ない。

一方で、宇宙線、突然writerの頭に隕石が、AtCoderがハッキングを食らう、など、そもそもコンテストが台無しになる可能性がそもそもそれなりにある。

典型ミス: 「想定解は(ハチャメチャ低い確率)でp < 0.5」 を言うには追加の仮定が必要

P(4000ケースAC | ($p \geq 0.5$))とP(($p \geq 0.5$) | 4000ケースAC)は違うという話。例えば、

  • こういう問題を作る
  • $p=0.01$の解法を作って5000ケース入れる
  • 4000ケースACしたら出題

これをとんでもない回数、例えば$100^{5000}$回行うと、このような問題が大量に出題され、そしてその全てが$p = 0.01$の大嘘解法となる。

もちろんこれは$100^{5000}$回の試行を要求している時点で非現実的である。全く効果のない薬をたくさん作ってテストし続けるとどれかは効果がありそうな結果が出てくる、というのと似た話。

  • 「4000/5000ケースACした」かつ「大量に試行していない」 => 「想定解は$p < 0.5$であると考えられる」
  • 「4000/5000ケースACした」かつ「$p$の事前分布が一様分布であると仮定する」 => 「想定解は(ハチャメチャ低い確率)で$p < 0.5$」

のように、追加で何らかの仮定が必要で、「4000/5000ケースACした」だけから「想定解は(ハチャメチャ低い確率)で$p < 0.5$」は数学的には言えない…ハズ

有意水準

$p$の事前分布はわからないだろうし、(今現在の)競プロ界隈のスタンス的に「writerが大量に試行していないから正しい」も広く受け入れられないと思う。なので$p$についてなにも言えなくて困る…と思いきや、こういう時のために有意水準というのがあるらしい(統計初心者)。つまり、 「$p < 0.5$なのに4000/5000ACするとんでもない確率を引いた」or「$p \geq 0.5$である」 ということならば言えるので、これで物事をやっていこうという話。

「確率 $q$ を引いた」or「X」が言えたとして、例えば $q$ が一生かけても引けないぐらいの確率(例: (1000年 / 試行にかかる時間) * q < 1)ならXは正しいと仮定して人生に問題はない?(例: 隕石は衝突しないと思い込んで行動しても人生に問題はない?)

薬の検定などと違い、テストケースをいくらでも簡単に増やせるのが強みで、それこそ$q \leq 2^{-256}$とかが達成できる問題も少なくないと思う。世の中は$2^{-256}$は起こらないということで回っているはず(例: 適当にEd25519の秘密鍵を当てられる確率が$2^{-256}$)なので、これなら大丈夫そう。

結論

  • 「$p < 0.5$なのに4000/5000ACするとんでもない確率を引いた」or「$p \geq 0.5$である」 ということならば言える。
  • こういう思想を(競プロで)許すか、許すとしてどのぐらいの確率からは個人次第。

自分の思想: とんでもない確率が$2^{-64}$とかならOK

思想についての余談

  • そもそも乱択に関する思想
    • 想定解の通る確率がケースごとに$p$以上が証明できていればOK派閥
    • 想定解の通る確率が全体で$p$以上が証明できていればOK派閥
      • ロリハ使うならケース数を公開しろ派閥がおそらくここ
    • 想定解は決定的(=確率1)アルゴリズムであるべき派閥
  • (1, 2番目の思想の場合)許せる確率$p$はいくつか
    • 全体で$\frac{1}{1000}$とか
    • 全体で$2^{-64}$とか(ハッシュの衝突など、世の中で0とされている確率ぐらい)
    • Full feedbackかにも依存しそう

自分の思想: ケースごとに$10^{-6}$ぐらいならええやろ派閥

ちょっと速いかもしれないローリングハッシュ

追記:速くなってませんでした!sorry https://x.com/yosupot/status/1689337328547016704?s=20

競技プログラミングではmod 261 - 1のローリングハッシュが安全性と速度のバランスが良く、広く使われています。 詳しくは https://qiita.com/keymoon/items/11fac5627672a6d6a9f6 などの記事が有用です。

このmod 261 - 1のmodintをより高速化することを試みます。先ほどの記事や、適当にライブラリを確認すると*1*2、乗算は以下の実装方法が広く使われています

using u64 = unsigned long long;
using u128 = unsigned __int128;

const u64 MOD = (1ULL << 61) - 1;

u64 mul(u64 a, u64 b) {
    u128 t = (u128)(a) * b;
    t = (t >> 61) + (t & MOD);

    return (t >= MOD) ? t - MOD : t;    
}

ここで、値をそのままではなく8倍して管理することを考えます。こうするとif (t >= MOD)相当の処理がoverflow checkになり、雰囲気的に良さそうな気がします。

u64 mul2(u64 a8, u64 b8) {
    u128 c = (u128)(a8) * b8;

    u64 x = (c >> 67 << 3), y = (c << 61 >> 64);

    u64 z;
    if (__builtin_uaddll_overflow(x, y, &z)) z -= MOD << 3;

    return z;
}

u64 mul(u64 a, u64 b) {
    u64 t = mul2(a * 8, b * 8) / 8;
    if (t == MOD) t = 0;
    return t;
}

実際に確認してみましょう。

godbolt.org

mul2のほうがめちゃくちゃすっきりしているのが確認できます。マジックナンバーがなく、本当に正しいのかこれという感じですが、読むと正しそうに思えます。

実際に O(N log2 N) SAを実装してみます

2割ほど早くなりました

注記

  • まだあんまり使ってないので、バグってるかも
  • そもそも従来のmodintのほうに改善の余地がありそう?tweet
  • 値を[0, MOD]で管理する都合上比較が汚くなってしまう これなんとかなるのかな? -> https://twitter.com/noshi91/status/1688130780718092288

Multiuni 2020 Day10 F

問題概要

問題

長さNのカッコ列 $a_1,a_2,\cdots,a_n$ が与えられる。これは正しく閉じているとは限らない。各文字は32bit非負整数の重み $b_1,b_2,\cdots,b_n$ を持っている。

カッコ列Sに対して、 $f(S)$ を次のように定義する

  • $S$ が ()を部分文字列として持っているかぎりそれを削除し続ける、その時の最終的な文字列。

$f(S)$はどのindexの文字を残すかまで含めて一意に定まることに注意。

次の $Q$ 個のクエリを処理。

  • 1 x y: $b_x$を $y$ に変更する(問題文には $a_x \to 1 - a_x$ もすると書かれているが、これは嘘)
  • 2 l r: $f(S[l..r])$ の重みを $c_1, c_2, \cdots, c_k$ としたときの、$\max (c_1, c_2, \cdots, c_k), \mathrm{nand}(...\mathrm{nand}(\mathrm{nand}(2^{32}-1, c_1), c_2), ..., c_k)$を求める。この2つをxorしたものを出力する。
  • 3 l r: $l..r$文字目と$r+1..n$文字目が入れ替わるようにswapする

制約

  • $N \leq 2 \times 10^{6}$
  • $Q \leq 2 \times 10^{5}$

解法

クエリ2で求めるmax / nandは共にモノイドの演算として考えることができる。

カッコ列から () を取り除き続けたときの最終的な文字列の長さが平衡二分木に乗るのはそこそこ有名。最終的な文字列は ))..)(..((という形になるので、 )( の個数を x, y とすると、 op((lx, ly), (rx, ry)) = (lx + rx - min(ly, rx), ly + ry - min(ly, ry) のような演算ができる。

今回の問題で同様のことをすると、ノードごとに

  • ))..)の長さ ln
  • ((..(の長さ rn
  • ))..)の重みの総和 lval
  • ((..(の重みの総和 rval

を持たせたくなるけど、これだとうまくノードがマージできない。

ここで、次の3つのパターンに限ればmerge可能なことに注目する。

  • l->rn = 0
  • r->ln = 0
  • l->rn = r->ln

「すべての葉以外のノードがこの3つの条件のいずれかを満たす」ような、葉に値を持たせる平衡二分木を管理する。

この追加条件を保つようにsplay treeを改造する。

方針としては、次のクエリを実装する。

  • lsplit(node, k): ノードを2つに分割する。左のノードに対応する文字列 $S$ について、$f(S)$ は ))..) ($k$ 文字)。
  • rsplit(node, k): ノードを2つに分割する。右のノードに対応する文字列 $S$ について、$f(S)$ は ((..(($k$ 文字)。

これらの関数が実装できると、通常の平衡二分木のようにmergeが実装できる。mergeの引数が上記の3条件を満たさない場合でも、lsplitかrsplitを呼ぶことでmerge可能な形に変形できる。

lsplitやrsplitは、このmerge関数が実装できれば実装できる。つまり相互再帰みたいな感じになる。

計算量

何もかもが謎

  • 手元で適当に試すと $O(\log N) / \mathrm{query}$ っぽい挙動をするが、はたして…
  • 実はもう少し違う方針で $O(\log ^2 N) / \mathrm{query}$ は達成できるが、これはTLEした
  • writer解も謎のsplay treeっぽいことをしていた、editorialがないのでこれも計算量は謎

コード

#include <cstdio>
#include <cassert>
#include <memory>
#include <algorithm>
#include <vector>

using namespace std;
using uint = unsigned int;

struct Monoid {
    uint mx, zero, one;
    Monoid() {
        mx = 0;
        zero = 0;
        one = -1;
    }
    Monoid(uint x) {
        mx = x;
        zero = -1;
        one = ~x;
    }
    uint eval() { return mx ^ one; }
};
Monoid operator+(const Monoid& l, const Monoid& r) {
    Monoid m;
    m.mx = max(l.mx, r.mx);
    m.zero = (l.zero & r.one) | (~l.zero & r.zero);
    m.one = (l.one & r.one) | (~l.one & r.zero);
    return m;
}

struct Node;
using NP = unique_ptr<Node>;

struct Node {
    NP l = nullptr, r = nullptr;
    int sz = -1;

    int ln, rn;
    Monoid lval, rval;

    Node() {}

    // leaf node, true='(', false=')'
    Node(bool type, uint x) : sz(1) {
        if (!type) {
            ln = 1;
            rn = 0;
            lval = Monoid(x);
            rval = Monoid();
        } else {
            ln = 0;
            rn = 1;
            lval = Monoid();
            rval = Monoid(x);
        }
    }
    // non leaf node
    Node(NP _l, NP _r) : l(move(_l)), r(move(_r)), sz(l->sz + r->sz) {
        assert(l && r);
        if (l->rn == r->ln) {
            ln = l->ln;
            rn = r->rn;
            lval = l->lval;
            rval = r->rval;
        } else if (l->rn == 0) {
            ln = l->ln + r->ln;
            rn = r->rn;
            lval = l->lval + r->lval;
            rval = r->rval;
        } else if (r->ln == 0) {
            ln = l->ln;
            rn = l->rn + r->rn;
            lval = l->lval;
            rval = l->rval + r->rval;
        } else {
            assert(false);
        }
    }
};

pair<NP, NP> lsplit(NP x, int k);
pair<NP, NP> rsplit(NP x, int k);

NP merge(NP l, NP r) {
    if (!l) return r;
    if (!r) return l;
    if (l->rn == 0 || r->ln == 0 || l->rn == r->ln) {
        return NP(new Node(move(l), move(r)));
    }

    if (l->rn < r->ln) {
        auto u = lsplit(move(r), l->rn);
        return NP(
            new Node(NP(new Node(move(l), move(u.first))), move(u.second)));
    } else {
        auto u = rsplit(move(l), r->ln);
        return NP(
            new Node(move(u.first), NP(new Node(move(u.second), move(r)))));
    }
}
template<class F>
pair<NP, NP> split2(NP x, F f) {
    int type = f(x);
    if (type == 0) {
        return {move(x->l), move(x->r)};
    }
    if (type == -1) {
        int type2 = f(x->l);
        if (type2 == 0) {
            return {move(x->l->l), merge(move(x->l->r), move(x->r))};
        }
        if (type2 == -1) {
            // zig-zig
            auto u = split2(move(x->l->l), f);
            return {move(u.first),
                    merge(move(u.second), merge(move(x->l->r), move(x->r)))};
        } else {
            // zig-zag
            auto u = split2(move(x->l->r), f);
            return {merge(move(x->l->l), move(u.first)),
                    merge(move(u.second), move(x->r))};
        }
    } else {
        int type2 = f(x->r);
        if (type2 == 0) {
            return {merge(move(x->l), move(x->r->l)), move(x->r->r)};
        }
        if (type2 == 1) {
            // zig-zig
            auto u = split2(move(x->r->r), f);
            return {merge(merge(move(x->l), move(x->r->l)), move(u.first)),
                    move(u.second)};
        } else {
            // zig-zag
            auto u = split2(move(x->r->l), f);
            return {merge(move(x->l), move(u.first)),
                    merge(move(u.second), move(x->r->r))};
        }
    }
}

pair<NP, NP> lsplit(NP x, int k) {
    assert(0 <= k && k <= x->ln);
    if (k == 0) {
        return {nullptr, move(x)};
    } else if (k == x->ln) {
        return {move(x), nullptr};
    }

    return split2(move(x), [&](const NP& n) {
        assert(0 < k && k < n->ln);
        int lsz = n->l->ln;
        if (lsz == k) return 0;
        if (k < lsz) return -1;
        k -= lsz;
        return 1;
    });
}

pair<NP, NP> rsplit(NP x, int k) {
    assert(0 <= k && k <= x->rn);
    if (k == 0) {
        return {move(x), nullptr};
    } else if (k == x->rn) {
        return {nullptr, move(x)};
    }

    return split2(move(x), [&](const NP& n) {
        assert(0 < k && k < n->rn);
        int rsz = n->r->rn;
        if (rsz == k) return 0;
        if (k < rsz) return 1;
        k -= rsz;
        return -1;
    });
}

pair<NP, NP> split(NP x, int k) {
    assert(0 <= k && k <= x->sz);
    if (k == 0) {
        return {nullptr, move(x)};
    } else if (k == x->sz) {
        return {move(x), nullptr};
    }

    return split2(move(x), [&](const NP& n) {
        assert(0 < k && k < n->sz);
        int lsz = n->l->sz;
        if (lsz == k) return 0;
        if (k < lsz) return -1;
        k -= lsz;
        return 1;
    });
}

int main() {
    int n, q;
    scanf("%d %d", &n, &q);

    vector<int> a(n);
    vector<uint> b(n);
    for (int i = 0; i < n; i++) {
        scanf("%d %d", &(a[i]), &(b[i]));
    }

    auto build = [&](auto self, int l, int r) -> NP {
        if (l + 1 == r) {
            return NP(new Node(a[l] == 1, b[l]));
        }
        int mid = (l + r) / 2;
        return merge(self(self, l, mid), self(self, mid, r));
    };
    NP tr = build(build, 0, n);

    for (int ph = 0; ph < q; ph++) {
        int ty, l, r;
        scanf("%d %d %d", &ty, &l, &r);
        l--;

        if (ty == 1) {
            auto t0 = split(move(tr), l + 1);
            auto t1 = split(move(t0.first), l);

            assert(t1.second->sz == 1);

            *t1.second = Node(t1.second->rn == 1, r);
            tr = merge(merge(move(t1.first), move(t1.second)), move(t0.second));
        } else if (ty == 2) {
            auto t0 = split(move(tr), r);
            auto t1 = split(move(t0.first), l);

            auto val = t1.second->lval + t1.second->rval;

            printf("%u\n", val.eval());

            tr = merge(merge(move(t1.first), move(t1.second)), move(t0.second));
        } else {
            auto t0 = split(move(tr), r);
            auto t1 = split(move(t0.first), l);

            tr = merge(merge(move(t1.first), move(t0.second)), move(t1.second));
        }
    }
}

Suffix Automaton

概要

文字列 $S$ のSuffix Automatonとは、ざっくりいうととても性質のいいDFA(決定性オートマトン)である。一番代表的な性質は次のとおりである。

  • $S$ の部分文字列全て、またそれらのみを受理する。
  • 頂点数,辺数が $O(|S|)$、より正確には $|V| \leq (2|S| - 1), |E| \leq (3|S| - 4)$
  • $O(|S|)$ 時間で構築可能

このようなオートマトンが存在するということがまず非自明なのだが、これらに加えて、更に様々な良い性質がある。

構築

構成

最初に、どのようなオートマトンを作るのか(および、存在性の証明)を示す。結論から言うと、Suffix Automatonは$\mathrm{rev}(S)$のCompressed Suffix Treeから機械的に作ることが出来る。

例えば、$S = "babcc"$、つまり $\mathrm{rev}(S) = "ccbab"$ の場合、Suffix Treeは次のようになる。黒色のノードは、そのノードを終端とする Suffix が存在することを表す。

f:id:yosupo:20210131160146j:plain

「次数 $1$ かつ白色のノード」を子のノードとマージしたものがCompressed Suffix Treeである。これと、$S$ の Suffix Automaton が次のようになる。

f:id:yosupo:20210131160121j:plain

ここで、左右のノードは一対一対応している。実際に、ノード7について、受理する文字列の集合が(revすると)全く同一である。他のノードに対しても同じ性質が成り立つ。

このようなAutomatonが必ず存在することが、次の定理によりわかる。証明は容易なので略する。

  • Compressed Suffix Treeの(根以外の)任意のノードについて、受理する文字列たちから最初の文字を削除した集合は、いくつかのノードの受理する文字列たちの直和で表せる。

実際に、Compressed Suffix Treeの各ノードについて、上記の直和のノードたちから最初の文字で遷移を貼るだけでSuffix Automatonが構築できる。

頂点数 / 辺数

頂点数は明らかに線形である(一般に、$n$ 個の文字列からパトリシアを作ると頂点数は$O(n)$になる)。

辺数はまずSuffix Automatonから(有向)全域木を取る。当然これの辺数はノード数 - 1である。全域木に含まれない辺それぞれについて、(全域木のパス + 含まれない辺)から生成される文字列たちを考える。すると

  • 全て $S$ の部分文字列
  • 文字列の間に、「片方が片方のprefix」という関係はない

が成立するので、文字列の個数 = 全域木に含まれない辺数 は高々 $|S|$

アルゴリズム

Compressed Suffix Tree / Suffix Automaton を対応を意識しながら並列で構築する。KMPアルゴリズムのように、$S$ の後ろ($\mathrm{rev}(S)$の先頭)に一文字ずつ文字を追加していく。これは Compressed Suffix Treeでは、(新しい)文字列全体をSuffix Treeに追加することに対応する。

Compressed Suffix Tree / Suffix Automatonの情報として、次の情報を持つ。

  • $\text{next}(\text{node}, \text{char})$ : Suffix Automatonの辺
  • $\mathrm{link}(\mathrm{node})$ : Compressed Suffix Treeでの親ノード
  • $\mathrm{len}(\mathrm{code})$ : nodeが受理する最長の文字列の長さ(=Suffix Treeでの一番下のノードの深さ)
  • $\mathrm{last}$ : $S$ 全体を入れた時に受理するノード

lenだけ不自然に感じるが、構築で必要になる。

構築においての本質は、Suffix Automatonの次の性質である。

  • (非空の文字列) $S$ の最後の文字を削除した文字列を $S'$ とする。Compressed Suffix Treeでの $S'$ / $S$ のパス上のノードを $n'1, n'2, \cdots, n'l$ / $n_1, n_2, \cdots, n_m$ とする($n'1 = n_1 = \mathrm{根}$)。このとき、$\mathrm{根}, \mathrm{next}(n'1, x), \mathrm{next}(n'2, x), \cdots, \mathrm{next}(n'_l, x)$ をランレングス圧縮したものが、$n_1, n_2, \cdots, n_m$ となる。

構築では、新しい文字を後ろに追加した後この性質を満たすように様々な値をいじればよい。

実際には新しいノードを作り、 $\mathrm{last}$ からlinkをたどっていく。そして

  • $\mathrm{next}(n, x)$ が存在しない間 $\to$ 新しく作ったノードに$\mathrm{next}(n, x)$を張っていく
  • $\mathrm{next}(n, x)$ が存在する場合。ノード $m = \mathrm{next}(n, x)$を分割(Clone)しないといけない場合がある。この判定に $\mathrm{len}$ を使う。$\mathrm{len}(n) + 1 = \mathrm{len}(m)$ ならば Cloneする必要がない。
    • 分割するためには、新規ノードを追加する。これはSuffix Treeで根から近い側に対応する。更に$m = \mathrm{next}(n, x)$ の間linkを辿り、nextを新規ノードに張り替える必要があることに注意。

構築の計算量について考える。頂点数、辺数が線形であるので大部分は大丈夫だが、唯一怪しいのは「ノードをCloneした後にnextを張り替える」部分である。実際には、「Compress Suffix Treeでの $\mathrm{last}$ ノードの高さ」をポテンシャルとすることで抑えることが出来る。

性質

その他の性質を列挙する

  • Suffix Automatonは「Sの部分列、またそれのみを受理する」DFAの中で頂点数が最小(らしい)
  • $S$ のSuffixと、cloneされてないノードたちは1対1対応する。

拡張

Compressed Suffix Treeが自然に複数の文字列に対応できるように、Suffix Automatonも対応できる。

lastを初期化して文字列追加、を同じ Suffix Automaton に繰り返せば良い。ここで、新しいノードが生まれない場合もあることに注意。

問題例

$S$ の部分文字列の種類数

  • DAGのパスの総数となるので、DP出来る
  • 各頂点 $n$ について$\mathrm{len}(n) - \mathrm{len}(\mathrm{link}(n))$ を足すだけでもよい(これはCompressed Suffix Treeでのパスの長さ=受理する文字列数に対応する)。

$S, T$の共通部分列の種類数

  • $S$, $T$, $S$ + "$" + $T$ の部分列の個数から計算できる。
  • ${ S, T }$ から Suffix Automaton を作ることで直接計算できる。$S$ のSuffix, $T$ のSuffixのどちらからも到達可能なノードのみについて$\mathrm{len}(n) - \mathrm{len}(\mathrm{link}(n))$を足せば良い

参考

使用例

AC code of Number of Substrings

struct SuffixAutomaton {
    struct Node {
        unordered_map<char, int> next;
        int link, len;
    };
    vector<Node> nodes;
    int last;

    SuffixAutomaton() {
        nodes.push_back({{}, -1, 0});
        last = 0;
    }

    void push(char c) {
        int new_node = int(nodes.size());
        nodes.push_back({{}, -1, nodes[last].len + 1});
        int p = last;
        while (p != -1 && nodes[p].next.find(c) == nodes[p].next.end()) {
            nodes[p].next[c] = new_node;
            p = nodes[p].link;
        }
        int q = (p == -1 ? 0 : nodes[p].next[c]);
        if (p == -1 || nodes[p].len + 1 == nodes[q].len) {
            nodes[new_node].link = q;
        } else {
            // clone node (q -> new_q)
            int new_q = int(nodes.size());
            nodes.push_back({nodes[q].next, nodes[q].link, nodes[p].len + 1});
            nodes[q].link = new_q;
            nodes[new_node].link = new_q;

            while (p != -1 && nodes[p].next[c] == q) {
                nodes[p].next[c] = new_q;
                p = nodes[p].link;
            }
        }
        last = new_node;
    }
};

int main() {
    string s;
    cin >> s;

    SuffixAutomaton sa;
    for (char c : s) {
        sa.push(c);
    }

    int m = int(sa.nodes.size());
    ll ans = 0;
    for (int i = 1; i < m; i++) {
        ans += sa.nodes[i].len - sa.nodes[sa.nodes[i].link].len;
    }

    cout << ans << endl;
    return 0;
}