bouzuya.hatenablog.com

ぼうずやのにっき

お腹が痛い / PAST#8 N を解いた

お腹が痛い……。


use std::{
    collections::{BTreeMap, BTreeSet},
    fmt::Display,
    ops::{Add, Sub},
};

use modint::ModInt1000000007 as ModInt;
use proconio::input;
use segtree::*;

use crate::internal_type_traits::Zero;

fn main() {
    input! {
        n: usize,
        a: [usize; n],
    }

    let map = a
        .iter()
        .copied()
        .collect::<BTreeSet<_>>()
        .into_iter()
        .enumerate()
        .fold(BTreeMap::new(), |mut m, (i, k)| {
            m.insert(k, i);
            m
        });

    #[derive(Clone, Copy)]
    struct M(ModInt);
    impl From<usize> for M {
        fn from(value: usize) -> Self {
            Self(ModInt::new(value))
        }
    }
    impl Add for M {
        type Output = M;
        fn add(self, rhs: Self) -> Self::Output {
            Self(self.0 + rhs.0)
        }
    }
    impl Display for M {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            self.0.fmt(f)
        }
    }
    impl Sub for M {
        type Output = M;
        fn sub(self, rhs: Self) -> Self::Output {
            Self(self.0 - rhs.0)
        }
    }
    impl Zero for M {
        fn zero() -> Self {
            M(ModInt::new(0))
        }
    }

    let mut zig = Segtree::<Additive<M>>::new(n);
    let mut zag = Segtree::<Additive<M>>::new(n);
    for a_i in a {
        let a_i = map[&a_i];
        let zig_p = zig.get(a_i);
        let zag_p = zag.get(a_i);
        zag.set(a_i, zag_p + zig.prod(0..a_i) + M::from(1));
        zig.set(a_i, zig_p + zag.prod(a_i + 1..n) + M::from(1));
    }

    let ans = zig.prod(0..) + zag.prod(0..) - M::from(2 * n);
    println!("{}", ans);
}

// modint
// segtree

今日のコミット。