全国統一プログラミング王決定戦本戦 D - Deforestation

本番は平面走査で通したけど、セグメント木で殴る解法にもちゃんと気付くべきだった。
D - Deforestation

問題

N本の竹がある。各竹の時刻0における高さは0で、毎秒1ずつ伸びる。
以下の行動をi = 1, ..., NのN回行うとき、総得点を求めよ。

時刻T[i]に、L[i]番目からR[i]番目までの竹を根元から切る。それらの竹の長さの合計を得点とする。切られた竹も、その後伸び続ける。

解法

数列の上の以下のようなクエリを捌ければよい。

区間和クエリ (L番目からR番目までの和を求める)
区間ゼロ初期化クエリ (L番目からR番目までを0にする)
区間加算クエリ (1番目からN番目までに1を加える)

これは以下のような遅延セグメント木で処理できる。

要素: i64
更新: (i64, i64)、(a, b)は要素xをax + bに変更するクエリを表す。

単位元および演算は以下の通り。

要素の単位元: 0
更新の単位元: (1, 0) (x |-> 1x + 0が恒等写像であるため)
要素の演算: (+)
更新の作用: ノードの高さ(根からの距離)をhとすると、xに(a,b)を作用させた結果はax + b * 2^h (理由: そのノードは2^h個の要素の和が格納されているため)
更新の演算: (a, b) * (c, d) = (ac, bc + d) (理由: yにx |-> ax + bとx |-> cx + dをこの順で適用すると、acy + bc + dになるため)

以上を遅延セグメント木に落とし込むとできる。これらの演算を書いた遅延セグメント木を書くのは大変なので、演算部分を除いて抽象化したものをライブラリにして、演算はtraitの実装という形で与えるのがよいだろう。

登場する典型

遅延セグメント木

実装上の注意点

更新の演算でバグらせないように気をつける。今回の場合、「ゼロクリア」と「区間加算」を同時に処理する構造を考えようとするとバグらせるに決まっているので、両方を含む代数構造を考えるとよい。

提出: Submission #4309587 - 全国統一プログラミング王決定戦本戦 (Rust)

/**
 * Lazy Segment Tree. This data structure is useful for fast folding and updating on intervals of an array
 * whose elements are elements of monoid T. Note that constructing this tree requires the identity
 * element of T and the operation of T. This is monomorphised, because of efficiency. T := i64, biop = max, upop = (+)
 * Reference: http://d.hatena.ne.jp/kyuridenamida/20121114/1352835261
 * Verified by https://codeforces.com/contest/1114/submission/49759034
 */
pub trait ActionRing {
    type T: Clone + Copy; // data
    type U: Clone + Copy + PartialEq + Eq; // action
    fn biop(Self::T, Self::T) -> Self::T;
    fn update(Self::T, Self::U, height: usize) -> Self::T;
    fn upop(Self::U, Self::U) -> Self::U;
    fn e() -> Self::T;
    fn upe() -> Self::U; // identity for upop
}
pub struct LazySegTree<R: ActionRing> {
    n: usize,
    dep: usize,
    dat: Vec<R::T>,
    lazy: Vec<R::U>,
}

impl<R: ActionRing> LazySegTree<R> {
    pub fn new(n_: usize) -> Self {
        let mut n = 1;
        let mut dep = 0;
        while n < n_ { n *= 2; dep += 1; } // n is a power of 2
        LazySegTree {
            n: n,
            dep: dep,
            dat: vec![R::e(); 2 * n - 1],
            lazy: vec![R::upe(); 2 * n - 1]
        }
    }
    #[inline]
    fn lazy_evaluate_node(&mut self, k: usize, height: usize) {
        if self.lazy[k] == R::upe() { return; }
        self.dat[k] = R::update(self.dat[k], self.lazy[k], height);
        if k < self.n - 1 {
            self.lazy[2 * k + 1] = R::upop(self.lazy[2 * k + 1], self.lazy[k]);
            self.lazy[2 * k + 2] = R::upop(self.lazy[2 * k + 2], self.lazy[k]);
        }
        self.lazy[k] = R::upe(); // identity for upop
    }
    #[inline]
    fn update_node(&mut self, k: usize) {
        self.dat[k] = R::biop(self.dat[2 * k + 1], self.dat[2 * k + 2]);
    }
    fn update_sub(&mut self, a: usize, b: usize, v: R::U, k: usize, height: usize, l: usize, r: usize) {
        self.lazy_evaluate_node(k, height);

        // [a,b) and  [l,r) intersects?
        if r <= a || b <= l {return;}
        if a <= l && r <= b {
            self.lazy[k] = R::upop(self.lazy[k], v);
            self.lazy_evaluate_node(k, height);
            return;
        }

        self.update_sub(a, b, v, 2 * k + 1, height - 1, l, (l + r) / 2);
        self.update_sub(a, b, v, 2 * k + 2, height - 1, (l + r) / 2, r);
        self.update_node(k);
    }
    /* ary[i] = upop(ary[i], v) for i in [a, b) (half-inclusive) */
    #[inline]
    pub fn update(&mut self, a: usize, b: usize, v: R::U) {
        let n = self.n;
        let dep = self.dep;
        self.update_sub(a, b, v, 0, dep, 0, n);
    }
    /* l,r are for simplicity */
    fn query_sub(&mut self, a: usize, b: usize, k: usize, height: usize, l: usize, r: usize) -> R::T {
        self.lazy_evaluate_node(k, height);

        // [a,b) and  [l,r) intersect?
        if r <= a || b <= l {return R::e();}
        if a <= l && r <= b {return self.dat[k];}
        let vl = self.query_sub(a, b, 2 * k + 1, height - 1, l, (l + r) / 2);
        let vr = self.query_sub(a, b, 2 * k + 2, height - 1, (l + r) / 2, r);
        self.update_node(k);
        R::biop(vl, vr)
    }
    /* [a, b) (note: half-inclusive) */
    #[inline]
    pub fn query(&mut self, a: usize, b: usize) -> R::T {
        let n = self.n;
        let dep = self.dep;
        self.query_sub(a, b, 0, dep, 0, n)
    }
}

struct Monoid {}
impl ActionRing for Monoid {
    type T = i64;
    type U = (i64, i64);
    fn e() -> i64 { 0 }
    fn upe() -> Self::U { (1, 0) }
    fn biop(a: i64, b: i64) -> i64 {
        a + b
    }
    fn update(a: Self::T, (b, c): Self::U, height: usize) -> i64 {
        a * b + (c << height)
    }
    fn upop((a, b): Self::U, (c, d): Self::U) -> Self::U {
        (a * c, b * c + d)
    }
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    input! {
        n: usize,
        tlr: [(i64, usize1, usize)],
    }
    let mut st = LazySegTree::<Monoid>::new(n);
    let mut last_t = 0;
    let mut tot = 0;
    for (t, l, r) in tlr {
        let diff = t - last_t;
        st.update(0, n, (1, diff));
        tot += st.query(l, r);
        st.update(l, r, (0, 0));
        last_t = t;
    }
    puts!("{}\n", tot);
}

まとめ

遅延セグメントツリーは実装が大変なので、うまい実装を書けるかどうかはライブラリ作成者の腕が問われる気がする。