D - Takahashi's Expectation Editorial by HBit

シミュレーションを高速化

テンションの変化を以下の連続した操作として言い換えることができます。

  1. 現在のテンションが \(P\) 以下のとき、テンションを \(A+B\) だけ増加させる。
  2. さらに、テンションを \(B\) だけ減少させる。

テンションをキー、そのテンションに対応するクエリ番号の集合を値とする順序付き連想配列を使用し、前からすべての質問を同時にシミュレートします。
操作 2 は実際に操作せず、連想配列のキーと実際のテンションの差分を別に管理します。
操作 1 は \(P\) が小さいことから愚直に追加削除を行うことができます。クエリ番号の集合のマージが必要になりますが、データ構造をマージする一般的なテクを使えば十分高速です。

全体の時間計算量は \(O((N\max(P)+Q)\log Q)\) です。


解答例 (Rust)

use proconio::{fastout, input};
use std::{collections::BTreeMap, mem::swap};

#[fastout]
pub fn main() {
    input!(n: usize, pab: [(usize, usize, usize); n], q: usize, x: [usize; q]);
    let mut map = BTreeMap::<usize, Vec<usize>>::new();
    for i in 0..q {
        map.entry(x[i]).or_default().push(i);
    }
    let mut acc = 0;
    let mut stack = vec![];
    for (p, a, b) in pab {
        while let Some((&key, _)) = map.range(..=p + acc).next() {
            let value = map.remove(&key).unwrap();
            stack.push((key.saturating_sub(acc) + acc + a + b, value));
        }
        while let Some((key, mut value)) = stack.pop() {
            map.entry(key)
                .and_modify(|e| {
                    if e.len() < value.len() {
                        swap(e, &mut value);
                    }
                    e.extend(&value);
                })
                .or_insert_with(|| value);
        }
        acc += b;
    }
    let mut ans = vec![0; q];
    for (key, value) in map {
        for i in value {
            ans[i] = key.saturating_sub(acc);
        }
    }
    for ans in ans {
        println!("{ans}");
    }
}

posted:
last update: