#918090 (Rust) No.447 ゆきこーだーの雨と雪 (2)

• 問題: No.447 ゆきこーだーの雨と雪 (2)
• ユーザー: ngtkana
• 提出日時: 2023-10-06 19:36:28
• 言語: Rust (1.83.0 + proconio)
• 結果: WA
• 実行時間: -
• コード長: 24,095 bytes
• コンパイル時間: 15,737 ms
• コンパイル使用メモリ: 377,732 KB
• 実行使用メモリ: 5,376 KB
• 最終ジャッジ日時: 2024-07-26 15:32:19
• 合計ジャッジ時間: 14,834 ms
• ジャッジサーバーID （参考情報）: judge1 / judge4

ファイルパターン	結果
sample	WA * 3
other	WA * 25

ソースコード

use crate::avl_tree::AvlTree;
use input::input;
use input::input_array;
use std::cmp::Reverse;
use std::collections::HashMap;
use std::iter::repeat;
use std::iter::repeat_with;

fn main() {
    let _ = input::<usize>();
    let allot = input::<String>()
        .split_whitespace()
        .map(|s| s.parse::<u64>().unwrap())
        .collect::<Vec<_>>();
    let (person_count, queries) = {
        let q = input::<usize>();
        let queries = repeat_with(|| input_array::<String, 2>())
            .take(q)
            .collect::<Vec<_>>();
        let mut names = queries
            .iter()
            .map(|query| query[0].as_str())
            .collect::<Vec<_>>();
        names.sort();
        names.dedup();
        let mut name_to_index = HashMap::new();
        for (i, &name) in names.iter().enumerate() {
            name_to_index.insert(name, i);
        }
        (
            names.len(),
            queries
                .iter()
                .map(|query| {
                    let name_index = name_to_index[query[0].as_str()];
                    match query[1].chars().next().unwrap() {
                        '?' => (name_index, Query::WhatRank),
                        pid => (name_index, Query::Solve((pid as u8 - b'A') as usize)),
                    }
                })
                .collect::<Vec<_>>(),
        )
    };

    let mut scores = vec![(0, Reverse(0)); person_count];
    let mut solved = vec![0; allot.len()];
    let mut sorted = repeat((0, Reverse(0)))
        .take(person_count)
        .collect::<AvlTree<_>>();
    for (time, &(person, query)) in queries.iter().enumerate() {
        match query {
            Query::WhatRank => {
                let rank = sorted.partition_point(|&node| scores[person] <= node);
                let ans = person_count - rank;
                println!("{}", ans);
            }
            Query::Solve(pid) => {
                sorted.remove(sorted.partition_point(|&node| scores[person] <= node));
                scores[person].0 += allot[pid] + allot[pid] * 5 / (5 + solved[pid]);
                scores[person].1 = Reverse(time);
                sorted.insert(
                    sorted.partition_point(|&node| scores[person] <= node),
                    scores[person],
                );
                solved[pid] += 1;
            }
        }
    }
}

#[derive(Debug, Clone, Copy)]
enum Query {
    Solve(usize),
    WhatRank,
}

// input {{{
#[allow(dead_code)]
mod input {
    use std::cell::Cell;
    use std::convert::TryFrom;
    use std::io::stdin;
    use std::io::BufRead;
    use std::io::BufReader;
    use std::io::Lines;
    use std::io::Stdin;
    use std::str::FromStr;
    use std::sync::Mutex;
    use std::sync::Once;
    type Server = Mutex<Lines<BufReader<Stdin>>>;
    static ONCE: Once = Once::new();
    pub struct Lazy(Cell<Option<Server>>);
    unsafe impl Sync for Lazy {}
    pub fn line() -> String {
        static SYNCER: Lazy = Lazy(Cell::new(None));
        ONCE.call_once(|| {
            SYNCER
                .0
                .set(Some(Mutex::new(BufReader::new(stdin()).lines())));
        });
        unsafe {
            (*SYNCER.0.as_ptr())
                .as_ref()
                .unwrap()
                .lock()
                .unwrap()
                .next()
                .unwrap()
                .unwrap()
        }
    }
    pub trait ForceFromStr: FromStr {
        fn force_from_str(s: &str) -> Self;
    }
    impl<T, E> ForceFromStr for T
    where
        T: FromStr<Err = E>,
        E: std::fmt::Debug,
    {
        fn force_from_str(s: &str) -> Self { s.parse().unwrap() }
    }
    pub fn input_array<T: ForceFromStr, const N: usize>() -> [T; N]
    where
        T: std::fmt::Debug,
    {
        <[_; N]>::try_from(input_vec()).unwrap()
    }
    pub fn input_vec<T: ForceFromStr>() -> Vec<T> {
        line()
            .split_whitespace()
            .map(T::force_from_str)
            .collect::<Vec<_>>()
    }
    pub fn input<T: ForceFromStr>() -> T { T::force_from_str(&line()) }
}
// }}}
// lg {{{
#[allow(dead_code)]
mod lg {
    use std::borrow::Borrow;
    use std::fmt::Debug;
    use std::marker::PhantomData;
    #[macro_export]
    macro_rules! lg {
        (@contents $head:expr $(, $tail:expr)*) => {{
            $crate::__lg_variable!($head);
            $(
                eprint!(",");
                $crate::__lg_variable!($tail);
            )*
            eprintln!();
        }};
        ($($expr:expr),* $(,)?) => {{
            eprint!("{}笶ｯ", line!());
            $crate::lg!(@contents $($expr),*)
        }};
    }
    #[doc(hidden)]
    #[macro_export]
    macro_rules! __lg_variable {
        ($value:expr) => {{
            match $value {
                head => {
                    eprint!(
                        " {} = {}",
                        stringify!($value),
                        $crate::lg::__quiet(format!("{:?}", &head))
                    );
                }
            }
        }};
    }
    #[doc(hidden)]
    pub fn __quiet(s: impl AsRef<str>) -> String {
        s.as_ref()
            .replace("18446744073709551615", "*") // u64
            .replace("9223372036854775807", "*") // i64
            .replace("-9223372036854775808", "*") // i64
            .replace("4294967295", "*") // u32
            .replace("2147483647", "*") // i32
            .replace("-2147483648", "*") // i32
            .replace("None", "*")
            .replace("Some", "")
    }
    #[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
    pub struct Table<T, Row, Storage> {
        __marker: PhantomData<(T, Row)>,
        storage: Storage,
        index_width: usize,
        column_width: usize,
        heading_newlines: usize,
    }
    pub fn table<T: Clone + Debug, Row: AsRef<[T]>, Storage: AsRef<[Row]>>(
        storage: Storage,
    ) -> Table<T, Row, Storage> {
        Table::new(storage)
    }
    impl<T, Row, Storage> Table<T, Row, Storage>
    where
        T: Clone + Debug,
        Row: AsRef<[T]>,
        Storage: AsRef<[Row]>,
    {
        pub fn new(storage: Storage) -> Self {
            Self {
                __marker: PhantomData,
                storage,
                column_width: 3,
                index_width: 2,
                heading_newlines: 1,
            }
        }

        pub fn index_width(&mut self, index_width: usize) -> &mut Self {
            self.index_width = index_width;
            self
        }

        pub fn column_width(&mut self, column_width: usize) -> &mut Self {
            self.column_width = column_width;
            self
        }

        pub fn heading_newlines(&mut self, heading_newlines: usize) -> &mut Self {
            self.heading_newlines = heading_newlines;
            self
        }
    }
    impl<T, Row, Storage> Debug for Table<T, Row, Storage>
    where
        T: Clone + Debug,
        Row: AsRef<[T]>,
        Storage: AsRef<[Row]>,
    {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let Self {
                __marker: _,
                ref storage,
                index_width,
                column_width,
                heading_newlines,
            } = *self;
            for _ in 0..heading_newlines {
                writeln!(f)?;
            }
            let ncols = storage.as_ref()[0].as_ref().len();
            write!(f, "\x1b[48;2;127;127;127;37m")?;
            write!(f, "{}|", " ".repeat(index_width))?;
            for j in 0..ncols {
                write!(f, "{j:column_width$}")?;
            }
            writeln!(f, "\x1b[0m")?;
            for (i, row) in storage.as_ref().iter().enumerate() {
                write!(f, "{:index_width$}|", i, index_width = index_width)?;
                for value in row.as_ref() {
                    write!(f, "{:>column_width$}", __quiet(format!("{:?}", value)),)?;
                }
                writeln!(f)?;
            }
            Ok(())
        }
    }
    pub fn bools<B, I>(iter: I) -> String
    where
        B: Borrow<bool>,
        I: IntoIterator<Item = B>,
    {
        format!(
            "[{}]",
            iter.into_iter()
                .map(|b| ['.', '#'][usize::from(*(b.borrow()))])
                .collect::<String>(),
        )
    }
}
// }}}
// avl_tree {{{
#[allow(dead_code)]
mod avl_tree {
    #![allow(clippy::unnecessary_box_returns)]
    use std::borrow::Borrow;
    use std::cmp::Ordering;
    use std::fmt::Debug;
    use std::hash::Hash;
    use std::iter::successors;
    use std::iter::FromIterator;
    use std::mem::swap;
    use std::ops::Index;
    #[derive(Clone)]
    pub struct AvlTree<T> {
        root: Option<Box<Node<T>>>,
    }
    impl<T> AvlTree<T> {
        pub fn new() -> Self { Self::default() }

        pub fn is_empty(&self) -> bool { self.root.is_none() }

        pub fn len(&self) -> usize { len(self.root.as_deref()) }

        pub fn push_back(&mut self, value: T) {
            self.append(&mut Self {
                root: Some(new(value)),
            })
        }

        pub fn push_front(&mut self, value: T) {
            let mut swp = Self {
                root: Some(new(value)),
            };
            swp.append(self);
            *self = swp;
        }

        pub fn pop_back(&mut self) -> Option<T> {
            let root = self.root.take()?;
            let last_index = root.len - 1;
            let (left, center, _right) = split_delete(root, last_index);
            self.root = left;
            Some(center.value)
        }

        pub fn pop_front(&mut self) -> Option<T> {
            let (_left, center, right) = split_delete(self.root.take()?, 0);
            self.root = right;
            Some(center.value)
        }

        pub fn back(&self) -> Option<&T> { self.get(self.len().checked_sub(1)?) }

        pub fn front(&self) -> Option<&T> { self.get(0) }

        pub fn back_mut(&mut self) -> Option<&mut T> { self.get_mut(self.len().checked_sub(1)?) }

        pub fn front_mut(&mut self) -> Option<&mut T> { self.get_mut(0) }

        pub fn append(&mut self, other: &mut Self) {
            self.root = merge(self.root.take(), other.root.take());
        }

        pub fn split_off(&mut self, index: usize) -> Self {
            assert!(index <= self.len());
            let (left, right) = split(self.root.take(), index);
            self.root = left;
            Self { root: right }
        }

        pub fn insert(&mut self, index: usize, value: T) {
            assert!(index <= self.len());
            let other = self.split_off(index);
            self.root = Some(merge_with_root(self.root.take(), new(value), other.root));
        }

        pub fn remove(&mut self, index: usize) -> Option<T> {
            if index < self.len() {
                let (left, center, right) = split_delete(self.root.take().unwrap(), index);
                self.root = merge(left, right);
                Some(center.value)
            } else {
                None
            }
        }

        pub fn get(&self, index: usize) -> Option<&T> {
            if index < self.len() {
                Some(&get(self.root.as_ref().unwrap(), index).value)
            } else {
                None
            }
        }

        pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
            if index < self.len() {
                Some(&mut get_mut(self.root.as_mut().unwrap(), index).value)
            } else {
                None
            }
        }

        pub fn binary_search_by(&self, f: impl FnMut(&T) -> Ordering) -> Result<usize, usize> {
            binary_search_by(self.root.as_deref(), f)
        }

        pub fn binary_search_by_key<B: Ord>(
            &self,
            b: &B,
            mut f: impl FnMut(&T) -> B,
        ) -> Result<usize, usize> {
            self.binary_search_by(|x| f(x).cmp(b))
        }

        pub fn binary_search<Q: Ord>(&self, value: &Q) -> Result<usize, usize>
        where
            T: Borrow<Q>,
        {
            self.binary_search_by(|x| x.borrow().cmp(value))
        }

        pub fn partition_point(&self, mut is_right: impl FnMut(&T) -> bool) -> usize {
            partition_point(self.root.as_deref(), |node| is_right(&node.value))
        }

        pub fn lower_bound<Q: Ord>(&self, value: &Q) -> usize
        where
            T: Borrow<Q>,
        {
            partition_point(self.root.as_deref(), |node| value <= node.value.borrow())
        }

        pub fn upper_bound<Q: Ord>(&self, value: &Q) -> usize
        where
            T: Borrow<Q>,
        {
            partition_point(self.root.as_deref(), |node| value < node.value.borrow())
        }

        pub fn iter(&self) -> Iter<'_, T> {
            Iter {
                stack: successors(self.root.as_deref(), |current| current.left.as_deref())
                    .collect(),
                rstack: successors(self.root.as_deref(), |current| current.right.as_deref())
                    .collect(),
            }
        }
    }
    impl<T> Default for AvlTree<T> {
        fn default() -> Self { Self { root: None } }
    }
    impl<T: PartialEq> PartialEq for AvlTree<T> {
        fn eq(&self, other: &Self) -> bool { self.iter().eq(other) }
    }
    impl<T: PartialEq, A> PartialEq<[A]> for AvlTree<T>
    where
        T: PartialEq<A>,
    {
        fn eq(&self, other: &[A]) -> bool { self.iter().eq(other) }
    }
    impl<T: Eq> Eq for AvlTree<T> {}
    impl<T: PartialOrd> PartialOrd for AvlTree<T> {
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> { self.iter().partial_cmp(other) }
    }
    impl<T: Ord> Ord for AvlTree<T> {
        fn cmp(&self, other: &Self) -> Ordering { self.iter().cmp(other) }
    }
    impl<T: Debug> Debug for AvlTree<T> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_list().entries(self).finish()
        }
    }
    impl<T: Hash> Hash for AvlTree<T> {
        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
            self.iter().for_each(|elm| elm.hash(state));
        }
    }
    impl<T> IntoIterator for AvlTree<T> {
        type IntoIter = IntoIter<T>;
        type Item = T;

        fn into_iter(self) -> Self::IntoIter {
            let mut stack = Vec::new();
            if let Some(mut current) = self.root {
                while let Some(next) = current.left.take() {
                    stack.push(current);
                    current = next;
                }
                stack.push(current);
            }
            IntoIter { stack }
        }
    }
    impl<'a, T> IntoIterator for &'a AvlTree<T> {
        type IntoIter = Iter<'a, T>;
        type Item = &'a T;

        fn into_iter(self) -> Self::IntoIter { self.iter() }
    }
    impl<T> Index<usize> for AvlTree<T> {
        type Output = T;

        fn index(&self, index: usize) -> &Self::Output { self.get(index).unwrap() }
    }
    impl<T> FromIterator<T> for AvlTree<T> {
        fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
            fn from_slice_of_nodes<T>(nodes: &mut [Option<Box<Node<T>>>]) -> Option<Box<Node<T>>> {
                if nodes.is_empty() {
                    None
                } else {
                    let i = nodes.len() / 2;
                    Some(merge_with_root(
                        from_slice_of_nodes(&mut nodes[..i]),
                        nodes[i].take().unwrap(),
                        from_slice_of_nodes(&mut nodes[i + 1..]),
                    ))
                }
            }
            Self {
                root: from_slice_of_nodes(
                    iter.into_iter()
                        .map(new)
                        .map(Some)
                        .collect::<Vec<_>>()
                        .as_mut_slice(),
                ),
            }
        }
    }
    pub struct Iter<'a, T> {
        stack: Vec<&'a Node<T>>,
        rstack: Vec<&'a Node<T>>,
    }
    impl<'a, T> Iterator for Iter<'a, T> {
        type Item = &'a T;

        fn next(&mut self) -> Option<Self::Item> {
            let current = self.stack.pop()?;
            self.stack
                .extend(successors(current.right.as_deref(), |node| {
                    node.left.as_deref()
                }));
            if std::ptr::eq(current, *self.rstack.last().unwrap()) {
                self.stack.clear();
                self.rstack.clear();
            }
            Some(&current.value)
        }
    }
    impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
        fn next_back(&mut self) -> Option<Self::Item> {
            let current = self.rstack.pop()?;
            self.rstack
                .extend(successors(current.left.as_deref(), |node| {
                    node.right.as_deref()
                }));
            if std::ptr::eq(current, *self.stack.last().unwrap()) {
                self.stack.clear();
                self.rstack.clear();
            }
            Some(&current.value)
        }
    }
    pub struct IntoIter<T> {
        stack: Vec<Box<Node<T>>>,
    }
    impl<T> Iterator for IntoIter<T> {
        type Item = T;

        fn next(&mut self) -> Option<Self::Item> {
            let mut current = self.stack.pop()?;
            if let Some(mut current) = current.right.take() {
                while let Some(next) = current.left.take() {
                    self.stack.push(current);
                    current = next;
                }
                self.stack.push(current);
            }
            Some(current.value)
        }
    }
    #[derive(Clone)]
    struct Node<T> {
        left: Option<Box<Self>>,
        right: Option<Box<Self>>,
        value: T,
        len: usize,
        ht: u8,
    }
    fn new<T>(value: T) -> Box<Node<T>> {
        Box::new(Node {
            left: None,
            right: None,
            value,
            len: 1,
            ht: 1,
        })
    }
    impl<T> Node<T> {
        fn update(&mut self) {
            self.len = len(self.left.as_deref()) + 1 + len(self.right.as_deref());
            self.ht = 1 + ht(self.left.as_deref()).max(ht(self.right.as_deref()));
        }
    }
    fn len<T>(tree: Option<&Node<T>>) -> usize { tree.as_ref().map_or(0, |node| node.len) }
    fn ht<T>(tree: Option<&Node<T>>) -> u8 { tree.as_ref().map_or(0, |node| node.ht) }
    fn balance<T>(node: &mut Box<Node<T>>) {
        fn rotate_left<T>(node: &mut Box<Node<T>>) {
            let mut x = node.left.take().unwrap();
            let y = x.right.take();
            swap(node, &mut x);
            x.left = y;
            x.update();
            node.right = Some(x);
            node.update();
        }
        fn rotate_right<T>(node: &mut Box<Node<T>>) {
            let mut x = node.right.take().unwrap();
            let y = x.left.take();
            swap(node, &mut x);
            x.right = y;
            x.update();
            node.left = Some(x);
            node.update();
        }
        if ht(node.left.as_deref()) > 1 + ht(node.right.as_deref()) {
            let left = node.left.as_mut().unwrap();
            if ht(left.left.as_deref()) < ht(left.right.as_deref()) {
                rotate_right(left);
            }
            rotate_left(node);
        } else if ht(node.left.as_deref()) + 1 < ht(node.right.as_deref()) {
            let right = node.right.as_mut().unwrap();
            if ht(right.left.as_deref()) > ht(right.right.as_deref()) {
                rotate_left(right);
            }
            rotate_right(node);
        } else {
            node.update();
        }
    }
    fn merge_with_root<T>(
        mut left: Option<Box<Node<T>>>,
        mut center: Box<Node<T>>,
        mut right: Option<Box<Node<T>>>,
    ) -> Box<Node<T>> {
        match ht(left.as_deref()).cmp(&ht(right.as_deref())) {
            Ordering::Less => {
                let mut root = right.take().unwrap();
                root.left = Some(merge_with_root(left, center, root.left.take()));
                balance(&mut root);
                root
            }
            Ordering::Equal => {
                center.left = left;
                center.right = right;
                center.update();
                center
            }
            Ordering::Greater => {
                let mut root = left.take().unwrap();
                root.right = Some(merge_with_root(root.right.take(), center, right));
                balance(&mut root);
                root
            }
        }
    }
    fn merge<T>(
        left: Option<Box<Node<T>>>,
        mut right: Option<Box<Node<T>>>,
    ) -> Option<Box<Node<T>>> {
        match right.take() {
            None => left,
            Some(right) => {
                let (_none, center, rhs) = split_delete(right, 0);
                Some(merge_with_root(left, center, rhs))
            }
        }
    }
    #[allow(clippy::type_complexity)]
    fn split_delete<T>(
        mut root: Box<Node<T>>,
        index: usize,
    ) -> (Option<Box<Node<T>>>, Box<Node<T>>, Option<Box<Node<T>>>) {
        debug_assert!((0..root.len).contains(&index));
        let left = root.left.take();
        let right = root.right.take();
        let lsize = len(left.as_deref());
        match lsize.cmp(&index) {
            Ordering::Less => {
                let mut res = split_delete(right.unwrap(), index - lsize - 1);
                res.0 = Some(merge_with_root(left, root, res.0));
                res
            }
            Ordering::Equal => (left, root, right),
            Ordering::Greater => {
                let mut res = split_delete(left.unwrap(), index);
                res.2 = Some(merge_with_root(res.2, root, right));
                res
            }
        }
    }
    #[allow(clippy::type_complexity)]
    fn split<T>(
        tree: Option<Box<Node<T>>>,
        index: usize,
    ) -> (Option<Box<Node<T>>>, Option<Box<Node<T>>>) {
        match tree {
            Some(root) => {
                if root.len == index {
                    (Some(root), None)
                } else {
                    let (left, center, right) = split_delete(root, index);
                    (left, Some(merge_with_root(None, center, right)))
                }
            }
            None => (None, None),
        }
    }
    fn binary_search_by<T>(
        tree: Option<&Node<T>>,
        mut f: impl FnMut(&T) -> Ordering,
    ) -> Result<usize, usize> {
        let node = match tree {
            None => return Err(0),
            Some(node) => node,
        };
        let lsize = len(node.left.as_deref());
        match f(&node.value) {
            Ordering::Less => binary_search_by(node.right.as_deref(), f)
                .map(|index| lsize + 1 + index)
                .map_err(|index| lsize + 1 + index),
            Ordering::Equal => Ok(lsize),
            Ordering::Greater => binary_search_by(node.left.as_deref(), f),
        }
    }
    fn partition_point<T>(
        tree: Option<&Node<T>>,
        mut is_right: impl FnMut(&Node<T>) -> bool,
    ) -> usize {
        let node = match tree {
            None => return 0,
            Some(node) => node,
        };
        let lsize = len(node.left.as_deref());
        if is_right(node) {
            partition_point(node.left.as_deref(), is_right)
        } else {
            lsize + 1 + partition_point(node.right.as_deref(), is_right)
        }
    }
    fn get<T>(node: &Node<T>, index: usize) -> &Node<T> {
        let lsize = len(node.left.as_deref());
        match lsize.cmp(&index) {
            Ordering::Less => get(node.right.as_ref().unwrap(), index - lsize - 1),
            Ordering::Equal => node,
            Ordering::Greater => get(node.left.as_ref().unwrap(), index),
        }
    }
    fn get_mut<T>(node: &mut Node<T>, index: usize) -> &mut Node<T> {
        let lsize = len(node.left.as_deref());
        match lsize.cmp(&index) {
            Ordering::Less => get_mut(node.right.as_mut().unwrap(), index - lsize - 1),
            Ordering::Equal => node,
            Ordering::Greater => get_mut(node.left.as_mut().unwrap(), index),
        }
    }
}
// }}}