はまやんはまやんはまやん

hamayanhamayan's blog

Subtree K-th Max [デンソークリエイトプログラミングコンテスト2022(AtCoder Beginner Contest 239) E]

https://atcoder.jp/contests/abc239/tasks/abc239_e

前提知識

解説

https://atcoder.jp/contests/abc239/submissions/29488895

データ構造でゴリ押して解けそうな感じもするし、Kの値が小さいことを利用して、楽に解くこともできそうではある。
根付き木の部分木に対する操作ではオイラーツアーが有効な場面が多々ある。

部分木に対する操作を配列の特定区間に対する操作に変換する

オイラーツアーを使うと部分木を配列の特定区間マッピングすることができる。
初めて聞く場合は何を言っているのか分からないと思うが、かなり便利なので検索して学習してこよう。
オイラーツアー - Google 検索
(ちょっと、ここに書くには余白が少なすぎる…)

オイラーツアーを使って部分木からインデックスを作成して、頂点Vに対する部分木に対応する区間が[L,R)に対応すると
分かったとする。すると、今回の問題を以下のように読み替えることができる。

頂点Vに対する部分木に含まれる数のうちK番目に大きい数は?

配列の区間[L, R)に含まれる数のうちK番目に大きい数は?

これで木であるということは無視してよくなるので、だいぶ今までの知見を活かしやすくなった。

特定区間のK番目に大きい数は?

これはWaveletMatrixというデータ構造を使用すると高速に取得できる。
しかし、結構高度なデータ構造なので、今回は、Kが小さいことを利用してセグメントツリーで解くことにする。
セグメントツリーを使えば最も大きい数というのは取得が可能である。
今回の問題ではK≦20と小さいので、

最も大きい数を取得する
→ 一旦それを除外する(正確には除外というか-1で更新する)
→ 最も大きい数を取得する(2番目に大きい数が得られる)
→ 一旦それも除外する
→ 最も大きい数を取得する(3番目に大きい数が得られる)

みたいなことを繰り返してK番目に大きい数を取得することにする。
取得した後は、除外した項目をセグメントツリーに戻してやることで取得操作の過程で壊れてしまった
セグメントツリーを復元していく。
こうすると、全体的な計算量はO(NKlogN)となり、K≦20,N≦105で2secだとちょっと不安な気もするが、
操作も比較的単純で定数倍も小さそうということもあり、まあ通るかという感じ。

advancedな話

  • wavelet matrixをやればストレートに区間問題パートが解けるので勉強してみるのも面白い
  • 自分の解法では消してundoみたいなことをしているが深堀したい場合は「永続化セグメントツリー」のような永続データ構造を調べてみても面白い
int N, Q, X[101010];
vector<int> E[101010];
int V[101010], K[101010];

int L[101010], R[101010];
int idx = 0;
void euler(int cu, int pa = -1) { // [L[v],R[v])
    L[cu] = idx; idx++;
    for (int to : E[cu]) if (to != pa) euler(to, cu);
    R[cu] = idx;
}

SegTree<1<<17> st;
void _main() {
    cin >> N >> Q;
    rep(i, 0, N) cin >> X[i];
    rep(i, 0, N - 1) {
        int a, b; cin >> a >> b;
        a--; b--;
        E[a].push_back(b);
        E[b].push_back(a);
    }
    rep(i, 0, Q) cin >> V[i] >> K[i], V[i]--;

    euler(0);

    rep(i, 0, N) st.update(L[i], { X[i], L[i] });
    
    rep(q, 0, Q) {
        int l = L[V[q]], r = R[V[q]];
        vector<pair<int, int>> buf;
        rep(k, 0, K[q] - 1) {
            auto p = st.get(l, r);
            buf.push_back(p);
            st.update(p.second, { -1, -1 });
        }

        auto p = st.get(l, r);
        printf("%d\n", p.first);
        fore(p, buf) st.update(p.second, p);
    }
}