空飛ぶ木造船

借り物ばかりの備忘録です。意味のあるものになると嬉しいです。

ABC253 C - Max - Min Queryについて

競技プログラミング

ABC253 C - Max - Min Queryについて、コンテスト中に自分が解いた方法とその理論的な解説、そして想定解法に沿った解答を確認したいと思います。

問題の要約

多重集合 Sに Q個のクエリが与えられるので、それらを処理する問題です。クエリは以下の三種類があります。

  • 1 x: Sに xを一つ追加
  • 2 x c: Sから xを \min(c, (S に含まれる x の個数 ))個削除する。
  • 3: ( Sの最大値)−( Sの最小値)を出力する。このクエリを処理するとき、 Sが空でないことが保証される。

コンテスト中に書いた解法と計算量

自分が考えた解き方は、辞書式配列mapと二つのヒープmax_heapとmin_heapを用意し、mapで要素の個数を管理しmax_heapで最大値、min_heapで最小値を取ることができるようにするというものでした。 実装は以下の通りになります。

use proconio::*;
use std::cmp::{min, Reverse};
use std::collections::{BinaryHeap, HashMap};

fn main() {
    input! {q: usize}

    let mut map = HashMap::new();
    let mut max_heap = BinaryHeap::new();
    let mut min_heap = BinaryHeap::new();

    for _ in 0..q {
        input! {kind: usize}

        match kind {
            1 => {
                input! {x: usize}
                *map.entry(x).or_insert(0) += 1;
                if map[&x] == 1 {
                    max_heap.push(x);
                    min_heap.push(Reverse(x));
                }
            }
            2 => {
                input! {
                    x: usize, c:usize
                }

                if map.contains_key(&x) {
                    *map.get_mut(&x).unwrap() -= min(c, map[&x]);
                }
            }
            3 => {
                let mut max;
                let mut min;
                loop {
                    let &m = max_heap.peek().unwrap();
                    if map[&m] > 0 {
                        max = m;
                        break;
                    } else {
                        max_heap.pop();
                    }
                }

                loop {
                    let &Reverse(m) = min_heap.peek().unwrap();
                    if map[&m] > 0 {
                        min = m;
                        break;
                    } else {
                        min_heap.pop();
                    }
                }

                println!("{}", max - min);
            }
            _ => return,
        }
    }
}

どうやらユーザー解説をみるとこの解法はmultisetがない場合の想定解のようなものだったようです。

ここで、一つ懸念点となるのが既に削除された値が最小値・最大値の候補になる場合です。この場合、候補の値が既にmap上で0個となっていれば(= 値が削除済み)、次の最大値・最小値をさらに確認し、それがmap上で1個以上の値を持っていればそれを最大値・最小値として出力します。この方法だと、例えば最大値の候補となる値が N個連続で軒並み削除済みの値であれば計算量が O(N \log N)になってしまいそうです。

この場合についての計算量をならし計算量(償却計算量)で考えます。まず一度の削除にかかる計算量はmap上の値を書き換えるだけなので O(1)です。そして、削除を N回行った後に最大値を出力しようとしてmax_heapから N回pop()を行うので、この計算量は O(N \log N)となりますが、これを連続した N回の操作全体で平均してやると、計算量は O(\log N)となります。

よって削除にかかる計算量も O(\log N)となります。上記の議論は [Python] heapqを応用して挿入/削除/最小値取得を効率的に行うを参考にしています。

想定解のRustによる実装

公式解説によれば、C++のmultisetクラスを利用するようなので、それと似たような機能をもつ構造体MultiSetを実装して解きます。 実装にあたっては記事Rustでmultisetもどき（ABC170 - E）を参考にHashSetやBTreeMapの要素の型をちょっと抽象化してみました。なんちゃって実装ですが、よかったら参考にしてください。

use proconio::*;
use std::cmp::min;
use std::collections::{BTreeSet, HashMap};
use std::hash::Hash;

struct MultiSet<T> {
    map: HashMap<T, usize>,
    set: BTreeSet<T>,
}

// 重複集合: 要素の重複を許し、要素が一意(= この場合は昇順)に並んだ集合
impl<T> MultiSet<T>
where
    T: Eq + Hash + Ord + Copy,
{
    fn new() -> Self {
        Self {
            map: HashMap::new(),
            set: BTreeSet::new(),
        }
    }

    // 重複許可挿入
    fn insert(&mut self, v: T) {
        *self.map.entry(v).or_insert(0) += 1;
        // setに要素がない場合は追加
        if self.map[&v] == 1 {
            self.set.insert(v);
        }
    }

    // 要素を一つ削除し、削除した要素を返す
    fn erase(&mut self, v: T) -> Option<T> {
        if !self.map.contains_key(&v) || self.map[&v] == 0 {
            return None;
        }

        *self.map.get_mut(&v).unwrap() -= 1;
        if self.map[&v] == 0 {
            self.set.take(&v);
        }
        Some(v)
    }

    // 要素vをc個削除。要素vがc個以下の場合はvの個数を0にする
    fn delete(&mut self, v: T, c: usize) {
        if !self.map.contains_key(&v) || self.map[&v] == 0 {
            return;
        }

        *self.map.get_mut(&v).unwrap() -= min(c, self.map[&v]);
        if self.map[&v] == 0 {
            self.set.take(&v);
        }
    }

    // 最大値を取得
    fn max(&self) -> Option<T> {
        if let Some(&v) = self.set.iter().last() {
            Some(v)
        } else {
            None
        }
    }

    // 最小値を取得
    fn min(&self) -> Option<T> {
        if let Some(&v) = self.set.iter().next() {
            Some(v)
        } else {
            None
        }
    }
}

fn main() {
    input! {q: usize}

    let mut multiset = MultiSet::new();

    for _ in 0..q {
        input! {kind: usize}

        match kind {
            1 => {
                input! {x: usize}
                multiset.insert(x);
            }
            2 => {
                input! {
                    x: usize, c:usize
                }
                multiset.delete(x, c);
            }
            3 => {
                println!("{}", multiset.max().unwrap() - multiset.min().unwrap());
            }
            _ => return,
        }
    }
}

参考文献