Submission #34093459

---

Source Code Expand

use std::cmp::Reverse;
use std::collections::BinaryHeap;

use proconio::input;

use nekolib::ds::{DecrementalUsizeSet, SkewHeap};

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

    let mut repr = DecrementalUsizeSet::new(n);

    let mut uq = BinaryHeap::new();
    let mut q: Vec<_> = (0..n).map(|_| SkewHeap::new()).collect();
    let mut size = vec![1; n];

    for i in 0..n - 1 {
        uq.push((Reverse(a[i] + a[i + 1]), (i, 1), (i + 1, 1)));
    }

    let mut res = 0;
    let mut a: Vec<_> = a.iter().map(|&ai| Some(ai)).collect();
    while let Some((Reverse(cost), (i, si), (j, sj))) = uq.pop() {
        if !(size[i] == si && size[j] == sj) {
            continue;
        }

        let gi = repr.less_equal(i).unwrap();
        let gj = repr.less_equal(j).unwrap();
        q[gi].pop();
        q[gj].pop();

        res += cost;
        size[i] += std::mem::replace(&mut size[j], 0);
        a[i] = Some(cost);
        a[j].take();
        q[gi].push((Reverse(cost), (i, size[i])));

        if let Some(r) = repr.greater(j) {
            if size[r] > 1 {
                let tmp = std::mem::take(&mut q[r]);
                q[j].meld(tmp);
                repr.remove(r);
            }
        }
        if repr.less_equal(gi) != repr.less_equal(j) {
            let tmp = std::mem::take(&mut q[j]);
            q[gi].meld(tmp);
            repr.remove(j);
        }

        if let Some(r) = repr.greater(gi) {
            if size[r] > 1 {
                let tmp = std::mem::take(&mut q[r]);
                q[gi].meld(tmp);
                repr.remove(r);
            }
        }
        if let Some(l) = repr.less(gi) {
            if size[l] > 1 {
                let tmp = std::mem::take(&mut q[gi]);
                q[l].meld(tmp);
                repr.remove(gi);
            }
        }

        let i = repr.less_equal(gi).unwrap();
        let fst = q[i].peek().unwrap().clone();
        if let Some(l) = repr.less(i) {
            let ncost = a[l].unwrap() + (fst.0).0;
            uq.push((Reverse(ncost), (l, 1), fst.1));
        }
        if let Some(r) = repr.greater(i) {
            let ncost = a[r].unwrap() + (fst.0).0;
            uq.push((Reverse(ncost), fst.1, (r, 1)));
        }

        if let (Some(l), Some(r)) = (repr.less(i), repr.greater(i)) {
            let ncost = a[l].unwrap() + a[r].unwrap();
            uq.push((Reverse(ncost), (l, 1), (r, 1)));
        }

        if q[i].len() > 1 {
            let fst = q[i].pop().unwrap();
            let snd = q[i].peek().unwrap().clone();
            q[i].push(fst); // back
            let ncost = (fst.0).0 + (snd.0).0;
            let left = fst.1.min(snd.1);
            let right = fst.1.max(snd.1);
            uq.push((Reverse(ncost), left, right));
        }
    }

    println!("{}", res);
}

/// This module is bundled automatically.
/// See <https://rsk0315.github.io/library-rs/nekolib/> for documentation.
pub mod nekolib {
    pub mod ds {
        pub mod decremental_usize_set {
            pub struct DecrementalUsizeSet {
                u: usize,
                len: usize,
                small: Vec<usize>,
                large: UnionFind,
            }
            const WORD_SIZE: usize = 0_usize.count_zeros() as usize;
            fn bsf(i: usize) -> usize { i.trailing_zeros() as usize }
            fn bsr(i: usize) -> usize {
                WORD_SIZE - 1 - i.leading_zeros() as usize
            }
            impl DecrementalUsizeSet {
                pub fn new(u: usize) -> Self {
                    let mut small = vec![!0_usize; u / WORD_SIZE + 1];
                    let mut large = UnionFind::new(small.len() + 1);
                    let div = u / WORD_SIZE;
                    let rem = u % WORD_SIZE;
                    small[div] = !(!0_usize << rem);
                    if rem == 0 {
                        large.unite(div, div + 1);
                    }
                    Self { u, len: u, small, large }
                }
                pub fn universe_len(&self) -> usize { self.u }
                pub fn len(&self) -> usize { self.len }
                pub fn is_empty(&self) -> bool { self.len == 0 }
                pub fn contains(&self, i: usize) -> bool {
                    let div = i / WORD_SIZE;
                    let rem = i % WORD_SIZE;
                    div < self.small.len() && self.small[div] >> rem & 1 != 0
                }
                pub fn less(&self, i: usize) -> Option<usize> {
                    if i == 0 { None } else { self.less_equal(i - 1) }
                }
                pub fn less_equal(&self, i: usize) -> Option<usize> {
                    let i = i.min(self.u - 1) + 1;
                    let div = i / WORD_SIZE;
                    let rem = i % WORD_SIZE;
                    let m = self.small[div] & !(!0_usize << rem);
                    if m != 0 {
                        return Some(div * WORD_SIZE + bsr(m));
                    }
                    let b = self.large.left(div);
                    if b == 0 {
                        None
                    } else {
                        Some((b - 1) * WORD_SIZE + bsr(self.small[b - 1]))
                    }
                }
                pub fn greater(&self, i: usize) -> Option<usize> {
                    if i == self.u { None } else { self.greater_equal(i + 1) }
                }
                pub fn greater_equal(&self, i: usize) -> Option<usize> {
                    let div = i / WORD_SIZE;
                    let rem = i % WORD_SIZE;
                    if div >= self.small.len() {
                        return None;
                    }
                    let m = self.small[div] & !0_usize << rem;
                    if m != 0 {
                        return Some(div * WORD_SIZE + bsf(m));
                    }
                    let b = self.large.right(div + 1);
                    if b == self.small.len() {
                        None
                    } else {
                        Some(b * WORD_SIZE + bsf(self.small[b]))
                    }
                }
                pub fn remove(&mut self, i: usize) -> bool {
                    let div = i / WORD_SIZE;
                    let rem = i % WORD_SIZE;
                    if div >= self.small.len()
                        || self.small[div] >> rem & 1 == 0
                    {
                        return false;
                    }
                    self.len -= 1;
                    self.small[div] &= !(1_usize << rem);
                    if self.small[div] == 0 {
                        self.large.unite(div, div + 1);
                    }
                    true
                }
            }
            #[derive(Clone, Copy)]
            enum Item {
                Parent(usize),
                Size(usize),
            }
            use std::cell::RefCell;
            use Item::{Parent, Size};
            struct UnionFind {
                buf: RefCell<Vec<Item>>,
                left: Vec<usize>,
                right: Vec<usize>,
            }
            impl UnionFind {
                pub fn new(n: usize) -> Self {
                    Self {
                        buf: RefCell::new(vec![Size(1); n]),
                        left: (0..n).collect(),
                        right: (0..n).collect(),
                    }
                }
                fn repr(&self, mut i: usize) -> usize {
                    let mut buf = self.buf.borrow_mut();
                    let res = {
                        let mut i = i;
                        while let Parent(ni) = buf[i] {
                            i = ni;
                        }
                        i
                    };
                    while let Parent(ni) = buf[i] {
                        buf[i] = Parent(res);
                        i = ni;
                    }
                    res
                }
                pub fn left(&self, i: usize) -> usize {
                    self.left[self.repr(i)]
                }
                pub fn right(&self, i: usize) -> usize {
                    self.right[self.repr(i)]
                }
                pub fn unite(&mut self, il: usize, ir: usize) {
                    let il = self.repr(il);
                    let ir = self.repr(ir);
                    let mut buf = self.buf.borrow_mut();
                    let (sl, sr) = match (buf[il], buf[ir]) {
                        (Size(sl), Size(sr)) => (sl, sr),
                        _ => unreachable!(),
                    };
                    if sl < sr {
                        buf[ir] = Size(sl + sr);
                        buf[il] = Parent(ir);
                        self.left[ir] = self.left[il];
                    } else {
                        buf[il] = Size(sl + sr);
                        buf[ir] = Parent(il);
                        self.right[il] = self.right[ir];
                    }
                }
            }
        }
        pub use self::decremental_usize_set::DecrementalUsizeSet;
        pub mod skew_heap {
            pub struct SkewHeap<T> {
                root: Link<T>,
                len: usize,
            }
            type Link<T> = Option<Box<Node<T>>>;
            struct Node<T> {
                elem: T,
                left: Link<T>,
                right: Link<T>,
            }
            impl<T> Node<T> {
                pub fn new(elem: T) -> Self {
                    Self { elem, left: None, right: None }
                }
            }
            impl<T: Ord> SkewHeap<T> {
                pub fn new() -> Self { SkewHeap { root: None, len: 0 } }
                pub fn len(&self) -> usize { self.len }
                pub fn is_empty(&self) -> bool { self.len == 0 }
                pub fn clear(&mut self) { while self.pop().is_some() {} }
                pub fn peek(&self) -> Option<&T> {
                    self.root.as_ref().map(|node| &node.elem)
                }
                pub fn peek_mut(&mut self) -> Option<PeekMut<'_, T>> {
                    if self.is_empty() {
                        None
                    } else {
                        Some(PeekMut { heap: self, tainted: false })
                    }
                }
                pub fn push(&mut self, elem: T) {
                    self.meld(Self {
                        root: Some(Box::new(Node::new(elem))),
                        len: 1,
                    });
                }
                pub fn pop(&mut self) -> Option<T> {
                    self.root.take().map(|node| {
                        self.len -= 1;
                        self.root = Self::meld_internal(node.left, node.right);
                        node.elem
                    })
                }
                pub fn meld(&mut self, other: Self) {
                    self.len += other.len;
                    self.root =
                        Self::meld_internal(self.root.take(), other.root);
                }
                fn meld_internal(left: Link<T>, right: Link<T>) -> Link<T> {
                    match (left, right) {
                        (left, None) => left,
                        (None, right) => right,
                        (Some(mut left), Some(right))
                            if left.elem >= right.elem =>
                        {
                            std::mem::swap(&mut left.left, &mut left.right);
                            left.left = Self::meld_internal(
                                left.left.take(),
                                Some(right),
                            );
                            Some(left)
                        }
                        (left, right) => Self::meld_internal(right, left),
                    }
                }
                fn fix_root(&mut self) {
                    let root = &self.root.as_ref().unwrap().elem;
                    let left = &self.root.as_ref().unwrap().left;
                    let right = &self.root.as_ref().unwrap().right;
                    if left
                        .as_ref()
                        .map(|left| &left.elem > root)
                        .unwrap_or(false)
                        || right
                            .as_ref()
                            .map(|right| &right.elem > root)
                            .unwrap_or(false)
                    {
                        let tmp = self.pop().unwrap();
                        self.push(tmp);
                    }
                }
            }
            impl<T> SkewHeap<T> {
                #[cfg(test)]
                fn push_with_dir(&mut self, elem: T, dir: &[bool]) {
                    let mut from = &mut self.root;
                    for &is_right in dir.iter().chain(std::iter::once(&false)) {
                        if let Some(v) = from {
                            from = if is_right {
                                &mut v.right
                            } else {
                                &mut v.left
                            };
                        } else if from.is_none() {
                            *from = Some(Box::new(Node::new(elem)));
                            return;
                        }
                    }
                }
                #[cfg(test)]
                fn in_order(&self) -> impl Iterator<Item = &T> {
                    fn dfs<'a, T>(v: &'a Link<T>, vec: &mut Vec<&'a T>) {
                        let v = match v {
                            Some(v) => v,
                            None => return,
                        };
                        dfs(&v.left, vec);
                        vec.push(&v.elem);
                        dfs(&v.right, vec);
                    }
                    let mut res = vec![];
                    dfs(&self.root, &mut res);
                    res.into_iter()
                }
                fn bfs_order(
                    &self,
                ) -> impl Iterator<Item = (&T, Option<&T>, Option<&T>)>
                {
                    use std::collections::VecDeque;
                    let mut res = vec![];
                    let mut q = VecDeque::new();
                    q.extend(self.root.as_ref());
                    while let Some(v) = q.pop_front() {
                        let left = v.left.as_ref();
                        let right = v.right.as_ref();
                        res.push((
                            &v.elem,
                            left.map(|o| &o.elem),
                            right.map(|o| &o.elem),
                        ));
                        q.extend(left.into_iter().chain(right));
                    }
                    res.into_iter()
                }
            }
            impl<T: Ord> Extend<T> for SkewHeap<T> {
                fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
                    for item in iter {
                        self.push(item);
                    }
                }
            }
            use std::iter::FromIterator;
            impl<T: Ord> FromIterator<T> for SkewHeap<T> {
                fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
                    let mut heap = Self::new();
                    heap.extend(iter);
                    heap
                }
            }
            pub struct IntoIter<T>(SkewHeap<T>);
            impl<T: Ord> IntoIterator for SkewHeap<T> {
                type IntoIter = IntoIter<T>;
                type Item = T;
                fn into_iter(self) -> Self::IntoIter { IntoIter(self) }
            }
            impl<T: Ord> Iterator for IntoIter<T> {
                type Item = T;
                fn next(&mut self) -> Option<Self::Item> { self.0.pop() }
                fn size_hint(&self) -> (usize, Option<usize>) {
                    (self.0.len, Some(self.0.len))
                }
            }
            impl<T: Ord> ExactSizeIterator for IntoIter<T> {
                fn len(&self) -> usize { self.0.len }
            }
            pub struct PeekMut<'a, T: 'a + Ord> {
                heap: &'a mut SkewHeap<T>,
                tainted: bool,
            }
            use std::ops::{Deref, DerefMut};
            impl<T: Ord> Drop for PeekMut<'_, T> {
                fn drop(&mut self) {
                    if self.tainted {
                        self.heap.fix_root();
                    }
                }
            }
            impl<T: Ord> Deref for PeekMut<'_, T> {
                type Target = T;
                fn deref(&self) -> &T { &self.heap.root.as_ref().unwrap().elem }
            }
            impl<T: Ord> DerefMut for PeekMut<'_, T> {
                fn deref_mut(&mut self) -> &mut T {
                    self.tainted = true;
                    &mut self.heap.root.as_mut().unwrap().elem
                }
            }
            impl<'a, T: Ord> PeekMut<'a, T> {
                pub fn pop(mut self) -> T {
                    let value = self.heap.pop().unwrap();
                    self.tainted = false;
                    value
                }
            }
            impl<T: Ord> Default for SkewHeap<T> {
                fn default() -> Self { Self::new() }
            }
            use std::fmt;
            impl<T: fmt::Debug> fmt::Debug for SkewHeap<T> {
                fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                    f.debug_list()
                        .entries(self.bfs_order().map(|(x, ..)| x))
                        .finish()
                }
            }
        }
        pub use self::skew_heap::SkewHeap;
    }
}

Submission Info

Judge Result