#652203 (Rust) No.1526 Sum of Mex 2

提出ソース
結果

問題	No.1526 Sum of Mex 2
ユーザー	Rheo Tommy
提出日時	2021-04-29 20:59:15
言語	Rust (1.77.0 + proconio)
結果	AC
実行時間	696 ms / 3,000 ms
コード長	18,537 bytes
コンパイル時間	15,591 ms
コンパイル使用メモリ	400,596 KB
実行使用メモリ	23,168 KB
最終ジャッジ日時	2024-11-08 21:22:37
合計ジャッジ時間	29,934 ms
ジャッジサーバーID （参考情報）	judge3 / judge2
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	1 ms 5,248 KB
testcase_01	AC	2 ms 5,248 KB
testcase_02	AC	2 ms 5,248 KB
testcase_03	AC	2 ms 5,248 KB
testcase_04	AC	1 ms 5,248 KB
testcase_05	AC	1 ms 5,248 KB
testcase_06	AC	2 ms 5,248 KB
testcase_07	AC	2 ms 5,248 KB
testcase_08	AC	2 ms 5,248 KB
testcase_09	AC	2 ms 5,248 KB
testcase_10	AC	2 ms 5,248 KB
testcase_11	AC	2 ms 5,248 KB
testcase_12	AC	2 ms 5,248 KB
testcase_13	AC	85 ms 5,248 KB
testcase_14	AC	99 ms 6,272 KB
testcase_15	AC	127 ms 6,784 KB
testcase_16	AC	667 ms 22,656 KB
testcase_17	AC	465 ms 19,580 KB
testcase_18	AC	44 ms 5,248 KB
testcase_19	AC	39 ms 5,248 KB
testcase_20	AC	275 ms 11,612 KB
testcase_21	AC	655 ms 22,100 KB
testcase_22	AC	359 ms 13,056 KB
testcase_23	AC	642 ms 22,528 KB
testcase_24	AC	690 ms 22,912 KB
testcase_25	AC	658 ms 22,272 KB
testcase_26	AC	667 ms 22,592 KB
testcase_27	AC	664 ms 22,888 KB
testcase_28	AC	696 ms 23,040 KB
testcase_29	AC	684 ms 22,872 KB
testcase_30	AC	692 ms 23,168 KB
testcase_31	AC	685 ms 23,168 KB
testcase_32	AC	694 ms 22,528 KB
testcase_33	AC	628 ms 22,096 KB
evil_largest	AC	2,088 ms 72,704 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
#![allow(unused_macros)]
#![allow(dead_code)]
#![allow(unused_imports)]

// # ファイル構成
// - use 宣言
// - lib モジュール
// - main 関数
// - basic モジュール
//
// 常に使うテンプレートライブラリは basic モジュール内にあります。
// 問題に応じて使うライブラリ lib モジュール内にコピペしています。
// ライブラリのコードはこちら → https://github.com/RheoTommy/at_coder
// Twitter はこちら → https://twitter.com/RheoTommy

use std::collections::*;
use std::io::{stdout, BufWriter, Write};

use crate::basic::*;
use crate::lib::*;

pub mod lib {
    pub trait LazyMonoid {
        type M: Clone;
        type L: Clone;
        fn e() -> Self::M;
        fn id() -> Self::L;
    }

    pub struct LazySegTree<M: LazyMonoid> {
        n: usize,
        height: usize,
        len: usize,
        data: Vec<(M::M, M::L)>,
        fold: Box<dyn FnMut(&M::M, &M::M) -> M::M>,
        eval: Box<dyn FnMut(&M::M, &M::L) -> M::M>,
        merge: Box<dyn FnMut(&M::L, &M::L) -> M::L>,
    }

    impl<M: LazyMonoid> std::fmt::Debug for LazySegTree<M>
        where
            <M as LazyMonoid>::M: std::fmt::Debug,
            <M as LazyMonoid>::L: std::fmt::Debug,
    {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            for i in 0..=self.height {
                for j in 0..1 << i {
                    write!(f, "{:?}   ", &self.data[(1 << i) + j])?;
                }
                writeln!(f)?;
            }
            write!(f, "")
        }
    }

    impl<M: LazyMonoid> LazySegTree<M> {
        pub fn new(
            n: usize,
            fold: Box<dyn FnMut(&M::M, &M::M) -> M::M>,
            eval: Box<dyn FnMut(&M::M, &M::L) -> M::M>,
            merge: Box<dyn FnMut(&M::L, &M::L) -> M::L>,
        ) -> Self {
            let len = n;
            let n = n.next_power_of_two();
            let height = n.trailing_zeros() as usize;
            Self {
                n,
                height,
                len,
                fold,
                eval,
                merge,
                data: vec![(M::e(), M::id()); 2 * n],
            }
        }
        pub fn from_slice(
            a: &[M::M],
            mut fold: Box<dyn FnMut(&M::M, &M::M) -> M::M>,
            eval: Box<dyn FnMut(&M::M, &M::L) -> M::M>,
            merge: Box<dyn FnMut(&M::L, &M::L) -> M::L>,
        ) -> Self {
            let n = a.len().next_power_of_two();
            let height = n.trailing_zeros() as usize;
            let mut data = vec![(M::e(), M::id()); 2 * n];
            for i in 0..a.len() {
                data[i + n].0 = a[i].clone();
            }
            for i in (1..n).rev() {
                data[i].0 = fold(&data[i * 2].0, &data[i * 2 + 1].0);
            }
            Self {
                n,
                height,
                len: a.len(),
                data,
                fold,
                eval,
                merge,
            }
        }
        fn apply(&mut self, x: usize, op: &M::L) {
            let node = &mut self.data[x];
            node.0 = (self.eval)(&node.0, op);
            node.1 = (self.merge)(&node.1, op);
        }
        fn propagate_at(&mut self, x: usize) {
            let op = std::mem::replace(&mut self.data[x].1, M::id());
            self.apply(2 * x, &op);
            self.apply(2 * x + 1, &op);
        }
        fn save_at(&mut self, x: usize) {
            self.data[x].0 = (self.fold)(&self.data[x * 2].0, &self.data[x * 2 + 1].0);
        }
        fn propagate(&mut self, l: usize, r: usize) {
            let l = l + self.n;
            let r = r + self.n;
            for i in (1..=self.height).rev() {
                if (l >> i) << i != l {
                    self.propagate_at(l >> i);
                }
                if (r >> i) << i != r {
                    self.propagate_at((r - 1) >> i);
                }
            }
        }
        fn save(&mut self, l: usize, r: usize) {
            let l = l + self.n;
            let r = r + self.n;
            for i in 1..=self.height {
                if (l >> i) << i != l {
                    self.save_at(l >> i);
                }
                if (r >> i) << i != r {
                    self.save_at((r - 1) >> i);
                }
            }
        }
        pub fn set(&mut self, l: usize, r: usize, op: M::L) {
            assert!(l <= r && r <= self.n);
            if l == r {
                return;
            }
            self.propagate(l, r);
            let mut x = l + self.n;
            let mut y = r + self.n;
            while x < y {
                if x & 1 == 1 {
                    self.apply(x, &op);
                    x += 1;
                }
                if y & 1 == 1 {
                    y -= 1;
                    self.apply(y, &op);
                }
                x >>= 1;
                y >>= 1;
            }
            self.save(l, r)
        }
        pub fn fold(&mut self, l: usize, r: usize) -> M::M {
            assert!(l <= r && r <= self.n);
            if l == r {
                return M::e();
            }
            self.propagate(l, r);
            let mut x = l + self.n;
            let mut y = r + self.n;
            let mut p = M::e();
            let mut q = M::e();
            while x < y {
                if x & 1 == 1 {
                    p = (self.fold)(&p, &self.data[x].0);
                    x += 1;
                }
                if y & 1 == 1 {
                    y -= 1;
                    q = (self.fold)(&self.data[y].0, &q);
                }
                x >>= 1;
                y >>= 1;
            }
            (self.fold)(&p, &q)
        }
        pub fn find_rightest(&mut self, l: usize, mut p: impl FnMut(&M::M) -> bool) -> usize {
            assert!(l < self.n);
            self.propagate(l, self.n);
            if p(&self.fold(l, self.n)) {
                return self.len;
            }
            let mut now = M::e();
            let mut res = l + self.n;
            while res + 1 != res.next_power_of_two() {
                if res & 1 == 1 {
                    let tmp = (self.fold)(&now, &self.data[res].0);
                    if !p(&tmp) {
                        break;
                    }
                    now = tmp;
                    res += 1;
                }
                res >>= 1;
            }
            while res < self.n {
                res <<= 1;
                let tmp = (self.fold)(&now, &self.data[res].0);
                if p(&tmp) {
                    now = tmp;
                    res += 1;
                }
            }
            (res - self.n).min(self.len)
        }
        pub fn find_leftest(&mut self, r: usize, mut p: impl FnMut(&M::M) -> bool) -> usize {
            assert!(r <= self.n);
            self.propagate(0, r);
            if p(&self.fold(0, r)) {
                return 0;
            }
            let mut now = M::e();
            let mut res = r + self.n;
            while res > 1 {
                if res & 1 == 1 {
                    res -= 1;
                    let tmp = (self.fold)(&self.data[res].0, &now);
                    if !p(&tmp) {
                        break;
                    }
                    now = tmp;
                }
                res >>= 1;
            }
            while res < self.n {
                res <<= 1;
                res += 1;
                let tmp = (self.fold)(&self.data[res].0, &now);
                if p(&tmp) {
                    res -= 1;
                    now = tmp;
                }
            }
            (res - self.n + 1).min(self.len)
        }
    }

    pub trait Monoid {
        type Item: Clone + std::fmt::Debug;
        fn op(a: &Self::Item, b: &Self::Item) -> Self::Item;
        fn id() -> Self::Item;
    }

    #[derive(Clone)]
    pub struct SegTree<M: Monoid> {
        n: usize,
        len: usize,
        data: Vec<M::Item>,
    }

    impl<M: Monoid> std::fmt::Debug for SegTree<M> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(f, "{:?}", &self.data[self.n - 1..])
        }
    }

    impl<M: Monoid> SegTree<M> {
        pub fn new(n: usize) -> Self {
            Self {
                len: n,
                n: n.next_power_of_two(),
                data: vec![M::id(); 2 * n.next_power_of_two() - 1],
            }
        }
        /// O(N) で a をもととする配列を構築
        pub fn from_slice(a: &[M::Item]) -> Self {
            let n = a.len().next_power_of_two();
            let mut data = vec![M::id(); 2 * n - 1];
            for i in 0..a.len() {
                data[i + n - 1] = a[i].clone();
            }
            for i in (0..n - 1).rev() {
                let left = &data[i * 2 + 1];
                let right = &data[i * 2 + 2];
                data[i] = M::op(left, right);
            }
            Self {
                n,
                data,
                len: a.len(),
            }
        }
        /// O(1) で data[i] を取得
        pub fn get(&self, i: usize) -> &M::Item {
            assert!(i < self.n);
            &self.data[i + self.n - 1]
        }
        /// O(logN) で data[i] <- x
        pub fn set(&mut self, i: usize, x: M::Item) {
            assert!(i < self.n);
            self.data[i + self.n - 1] = x;
            let mut j = i + self.n - 1;
            if j == 0 {
                return;
            }
            while {
                j = (j - 1) / 2;
                let left = &self.data[2 * j + 1];
                let right = &self.data[2 * j + 2];
                self.data[j] = M::op(left, right);
                j != 0
            } {}
        }
        /// O(logN) で [l,r) の区間和を返す
        pub fn fold(&self, mut l: usize, mut r: usize) -> M::Item {
            assert!(l <= r);
            assert!(r <= self.n);
            l += self.n;
            r += self.n;
            let mut left = M::id();
            let mut right = M::id();
            while l < r {
                if l & 1 == 1 {
                    left = M::op(&left, &self.data[l - 1]);
                    l += 1;
                }
                if r & 1 == 1 {
                    r -= 1;
                    right = M::op(&self.data[r - 1], &right);
                }
                l >>= 1;
                r >>= 1;
            }
            M::op(&left, &right)
        }
        /// O(logN) で、p(fold(l, res)) が真となるような最右端 res を返す
        pub fn find_rightest(&self, l: usize, mut p: impl FnMut(&M::Item) -> bool) -> usize {
            assert!(l < self.n);
            if p(&self.fold(l, self.n)) {
                return self.len;
            }
            let mut now = M::id();
            let mut res = l + self.n;
            while res + 1 != res.next_power_of_two() {
                if res & 1 == 1 {
                    let tmp = M::op(&now, &self.data[res - 1]);
                    if !p(&tmp) {
                        break;
                    }
                    now = tmp;
                    res += 1;
                }
                res >>= 1;
            }
            while res < self.n {
                res <<= 1;
                let tmp = M::op(&now, &self.data[res - 1]);
                if p(&tmp) {
                    now = tmp;
                    res += 1;
                }
            }
            (res - self.n).min(self.len)
        }
        /// O(logN) で、p(fold(res, r)) が真となるような最左端 res を返す
        pub fn find_leftest(&self, r: usize, mut p: impl FnMut(&M::Item) -> bool) -> usize {
            assert!(r <= self.n);
            if p(&self.fold(0, r)) {
                return 0;
            }
            let mut now = M::id();
            let mut res = r + self.n;
            while res > 1 {
                if res & 1 == 1 {
                    res -= 1;
                    let tmp = M::op(&self.data[res - 1], &now);
                    if !p(&tmp) {
                        break;
                    }
                    now = tmp;
                }
                res >>= 1;
            }
            while res < self.n {
                res <<= 1;
                res += 1;
                let tmp = M::op(&self.data[res - 1], &now);
                if p(&tmp) {
                    res -= 1;
                    now = tmp;
                }
            }
            (res - self.n + 1).min(self.len)
        }
    }

    pub struct Max;

    pub struct Min;

    pub struct Add;

    pub struct Mul;

    pub struct Xor;

    impl Monoid for Max {
        type Item = usize;
        fn id() -> Self::Item {
            0
        }
        fn op(a: &Self::Item, b: &Self::Item) -> Self::Item {
            *a.max(b)
        }
    }

    impl Monoid for Min {
        type Item = i64;
        fn id() -> Self::Item {
            (1 << 60) + (1 << 30)
        }
        fn op(a: &Self::Item, b: &Self::Item) -> Self::Item {
            *a.min(b)
        }
    }

    impl Monoid for Add {
        type Item = i64;
        fn id() -> Self::Item {
            0
        }
        fn op(a: &Self::Item, b: &Self::Item) -> Self::Item {
            a + b
        }
    }

    impl Monoid for Mul {
        type Item = i64;
        fn id() -> Self::Item {
            1
        }
        fn op(a: &Self::Item, b: &Self::Item) -> Self::Item {
            a * b
        }
    }

    impl Monoid for Xor {
        type Item = i64;
        fn id() -> Self::Item {
            0
        }
        fn op(a: &Self::Item, b: &Self::Item) -> Self::Item {
            a ^ b
        }
    }

    pub type MaxSegTree = SegTree<Max>;
    pub type MinSegTree = SegTree<Min>;
    pub type AddSegTree = SegTree<Add>;
    pub type MulSegTree = SegTree<Mul>;
    pub type XorSegTree = SegTree<Xor>;
}

struct M;

impl LazyMonoid for M {
    type M = (usize, usize);
    type L = Option<(usize, usize)>;

    fn e() -> Self::M {
        (0, 0)
    }

    fn id() -> Self::L {
        None
    }
}

fn main() {
    let mut io = IO::new();
    let n = io.next_usize();
    let a = io.next_vec::<usize>(n);
    let mut lst = LazySegTree::<M>::new(
        n + 2,
        Box::new(|a, b| (a.0 + b.0, a.1 + b.1)),
        Box::new(|a, b| if b.is_none() { *a } else { b.unwrap() }),
        Box::new(|a, b| match (a, b) {
            (_, Some(b)) => Some(*b),
            (a, _) => *a,
        }),
    );

    let mut index = vec![vec![n]; n + 1];
    index[0] = (0..=n).collect::<Vec<_>>();
    for i in 0..n {
        index[a[i]].push(i);
    }
    index.iter_mut().for_each(|vi| {
        vi.sort();
        vi.reverse();
    });

    let mut st = MaxSegTree::new(n + 1);
    for ai in 0..=n {
        st.set(ai, *index[ai].last().unwrap());
    }

    for ai in 0..=n {
        let i = st.fold(0, ai + 1);
        let (cnt, sum) = lst.fold(i, i + 1);
        lst.set(i, i + 1, Some((cnt + 1, sum + i)));
    }

    let mut ans = 0;
    for l in 0..n {
        ans += n * (n + 1) - lst.fold(0, n + 2).1;

        let ai = a[l];
        let now = st.fold(0, ai + 1);
        let r = st.find_rightest(0, |aj| *aj <= now);
        let len = r - ai;
        let (cnt, sum) = lst.fold(now, now + 1);
        lst.set(now, now + 1, Some((cnt - len, sum - len * now)));

        index[ai].pop();
        let next = *index[ai].last().unwrap().max(&st.fold(0, ai + 1));
        let amount = lst.fold((now + 1).min(next), next).0;
        lst.set((now + 1).min(next), next, Some((0, 0)));
        let (cnt, sum) = lst.fold(next, next + 1);
        lst.set(next, next + 1, Some((cnt + amount + len, sum + next * (len + amount))));
        st.set(ai, next);
        st.set(0, l + 1);

        lst.set(l, l + 1, Some((0, 0)));
        let (cnt, sum) = lst.fold(l + 1, l + 2);
        lst.set(l + 1, l + 2, Some((cnt + 1, sum + l + 1)));
    }

    io.println(ans);
}

pub mod basic {
    pub const U_INF: u64 = (1 << 60) + (1 << 30);
    pub const I_INF: i64 = (1 << 60) + (1 << 30);

    pub struct IO {
        iter: std::str::SplitAsciiWhitespace<'static>,
        buf: std::io::BufWriter<std::io::StdoutLock<'static>>,
    }

    impl Default for IO {
        fn default() -> Self {
            Self::new()
        }
    }

    impl IO {
        pub fn new() -> Self {
            use std::io::*;
            let mut input = String::new();
            std::io::stdin().read_to_string(&mut input).unwrap();
            let input = Box::leak(input.into_boxed_str());
            let out = Box::new(stdout());
            IO {
                iter: input.split_ascii_whitespace(),
                buf: BufWriter::new(Box::leak(out).lock()),
            }
        }
        pub fn next_str(&mut self) -> &str {
            self.iter.next().unwrap()
        }
        pub fn read<T: std::str::FromStr>(&mut self) -> T
            where
                <T as std::str::FromStr>::Err: std::fmt::Debug,
        {
            self.iter.next().unwrap().parse().unwrap()
        }
        pub fn next_usize(&mut self) -> usize {
            self.read()
        }
        pub fn next_uint(&mut self) -> u64 {
            self.read()
        }
        pub fn next_int(&mut self) -> i64 {
            self.read()
        }
        pub fn next_float(&mut self) -> f64 {
            self.read()
        }
        pub fn next_chars(&mut self) -> std::str::Chars {
            self.next_str().chars()
        }
        pub fn next_vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T>
            where
                <T as std::str::FromStr>::Err: std::fmt::Debug,
        {
            (0..n).map(|_| self.read()).collect::<Vec<_>>()
        }
        pub fn print<T: std::fmt::Display>(&mut self, t: T) {
            use std::io::Write;
            write!(self.buf, "{}", t).unwrap();
        }
        pub fn println<T: std::fmt::Display>(&mut self, t: T) {
            self.print(t);
            self.print("\n");
        }
        pub fn print_iter<T: std::fmt::Display, I: Iterator<Item=T>>(
            &mut self,
            mut iter: I,
            sep: &str,
        ) {
            if let Some(v) = iter.next() {
                self.print(v);
                for vi in iter {
                    self.print(sep);
                    self.print(vi);
                }
            }
            self.print("\n");
        }
        pub fn flush(&mut self) {
            use std::io::Write;
            self.buf.flush().unwrap();
        }
    }
}
yukicoder

結果

テストケース

ソースコード
