1851(E,rs): try solving “Nastya and Potions”

Signed-off-by: Matej Focko <me@mfocko.xyz>

1851/src/bin/e.rs

@@ -0,0 +1,311 @@
+#![allow(unused_imports)]
+
+// region ‹use›
+use self::data_structures::*;
+use self::input::*;
+use self::math::*;
+use self::output::*;
+use std::cmp::{max, min};
+use std::collections::HashMap;
+use std::collections::HashSet;
+// endregion ‹use›
+
+#[derive(PartialEq, Eq, Hash)]
+struct Way {
+    required: Vec<usize>,
+}
+
+impl Way {
+    fn estimated_cost(&self, costs: &[u64]) -> u64 {
+        if self.required.is_empty() {
+            return u64::MAX;
+        }
+
+        self.required.iter().map(|i| costs[i - 1]).sum()
+    }
+}
+
+fn find_costs(costs: &mut Vec<u64>, unlimited: Vec<usize>, ways: Vec<Way>) {
+    let mut is_final = vec![false; costs.len()];
+
+    // set cost of unlimited
+    for i in unlimited {
+        costs[i - 1] = 0;
+        is_final[i - 1] = true;
+    }
+
+    // set cost of inconstructible
+    let mut q: HashSet<usize> = HashSet::new();
+    for (i, way) in ways.iter().enumerate() {
+        if way.estimated_cost(costs) == u64::MAX {
+            is_final[i] = true;
+        } else {
+            q.insert(i);
+        }
+    }
+
+    while let Some(i) = q
+        .iter()
+        .filter(|i| !is_final[**i])
+        .min_by_key(|i| ways[**i].estimated_cost(costs))
+        .map(|i| *i)
+    {
+        costs[i] = min(costs[i], ways[i].estimated_cost(costs));
+        is_final[i] = true;
+        q.remove(&i);
+    }
+}
+
+fn solve(s: &mut Scanner) {
+    let n = s.next::<usize>();
+    let k = s.next::<usize>();
+
+    let mut costs = s.next_vec::<u64>(n);
+    let unlimited = s.next_vec::<usize>(k);
+
+    let ways = (0..n)
+        .map(|_| {
+            let m = s.next::<usize>();
+            Way {
+                required: s.next_vec::<usize>(m),
+            }
+        })
+        .collect();
+
+    find_costs(&mut costs, unlimited, ways);
+
+    let costs: Vec<String> = costs.iter().map(|c| format!("{}", c)).collect();
+    println!("{}", &costs.join(" "));
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn example_1() {
+        assert_eq!(1, 2);
+    }
+}
+
+// region runner
+const SINGLE_TEST: bool = false;
+fn main() {
+    let mut s = Scanner::new();
+
+    if SINGLE_TEST {
+        solve(&mut s)
+    } else {
+        let n = s.next::<usize>();
+        for _ in 0..n {
+            solve(&mut s)
+        }
+    }
+}
+// endregion runner
+
+#[allow(dead_code)]
+mod math {
+    const MOD: i64 = 1_000_000_007;
+
+    pub fn add(a: i64, b: i64) -> i64 {
+        (a + b) % MOD
+    }
+
+    pub fn sub(a: i64, b: i64) -> i64 {
+        ((a - b) % MOD + MOD) % MOD
+    }
+
+    pub fn mul(a: i64, b: i64) -> i64 {
+        (a * b) % MOD
+    }
+
+    pub fn exp(b: i64, e: i64) -> i64 {
+        if e == 0 {
+            return 1;
+        }
+
+        let half = exp(b, e / 2);
+        if e % 2 == 0 {
+            return mul(half, half);
+        }
+
+        mul(half, mul(half, b))
+    }
+
+    /// A trait implementing the unsigned bit shifts.
+    pub trait UnsignedShift {
+        fn unsigned_shl(self, n: u32) -> Self;
+        fn unsigned_shr(self, n: u32) -> Self;
+    }
+
+    /// A trait implementing the integer square root.
+    pub trait ISqrt {
+        fn isqrt(&self) -> Self
+        where
+            Self: Sized,
+        {
+            self.isqrt_checked()
+                .expect("cannot calculate square root of negative number")
+        }
+
+        fn isqrt_checked(&self) -> Option<Self>
+        where
+            Self: Sized;
+    }
+
+    macro_rules! math_traits_impl {
+        ($T:ty, $U: ty) => {
+            impl UnsignedShift for $T {
+                #[inline]
+                fn unsigned_shl(self, n: u32) -> Self {
+                    ((self as $U) << n) as $T
+                }
+
+                #[inline]
+                fn unsigned_shr(self, n: u32) -> Self {
+                    ((self as $U) >> n) as $T
+                }
+            }
+
+            impl ISqrt for $T {
+                #[inline]
+                fn isqrt_checked(&self) -> Option<Self> {
+                    use core::cmp::Ordering;
+                    match self.cmp(&<$T>::default()) {
+                        // Hopefully this will be stripped for unsigned numbers (impossible condition)
+                        Ordering::Less => return None,
+                        Ordering::Equal => return Some(<$T>::default()),
+                        _ => {}
+                    }
+
+                    // Compute bit, the largest power of 4 <= n
+                    let max_shift: u32 = <$T>::default().leading_zeros() - 1;
+                    let shift: u32 = (max_shift - self.leading_zeros()) & !1;
+                    let mut bit = <$T>::try_from(1).unwrap().unsigned_shl(shift);
+
+                    // Algorithm based on the implementation in:
+                    // https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Binary_numeral_system_(base_2)
+                    // Note that result/bit are logically unsigned (even if T is signed).
+                    let mut n = *self;
+                    let mut result = <$T>::default();
+                    while bit != <$T>::default() {
+                        if n >= (result + bit) {
+                            n -= result + bit;
+                            result = result.unsigned_shr(1) + bit;
+                        } else {
+                            result = result.unsigned_shr(1);
+                        }
+                        bit = bit.unsigned_shr(2);
+                    }
+                    Some(result)
+                }
+            }
+        };
+    }
+
+    math_traits_impl!(i8, u8);
+    math_traits_impl!(u8, u8);
+    math_traits_impl!(i16, u16);
+    math_traits_impl!(u16, u16);
+    math_traits_impl!(i32, u32);
+    math_traits_impl!(u32, u32);
+    math_traits_impl!(i64, u64);
+    math_traits_impl!(u64, u64);
+    math_traits_impl!(i128, u128);
+    math_traits_impl!(u128, u128);
+    math_traits_impl!(isize, usize);
+    math_traits_impl!(usize, usize);
+}
+
+#[allow(dead_code)]
+mod data_structures {
+    use std::cmp::{Ord, Reverse};
+    use std::collections::BinaryHeap;
+
+    #[derive(Debug)]
+    pub struct MinHeap<T> {
+        heap: BinaryHeap<Reverse<T>>,
+    }
+
+    impl<T: Ord> MinHeap<T> {
+        pub fn new() -> MinHeap<T> {
+            MinHeap {
+                heap: BinaryHeap::new(),
+            }
+        }
+
+        pub fn push(&mut self, item: T) {
+            self.heap.push(Reverse(item))
+        }
+
+        pub fn pop(&mut self) -> Option<T> {
+            self.heap.pop().map(|Reverse(x)| x)
+        }
+    }
+
+    impl<T: Ord> Default for MinHeap<T> {
+        fn default() -> Self {
+            Self::new()
+        }
+    }
+}
+
+#[allow(dead_code)]
+mod output {
+    pub fn yes() {
+        println!("YES");
+    }
+
+    pub fn no() {
+        println!("NO");
+    }
+
+    pub fn yesno(ans: bool) {
+        println!("{}", if ans { "YES" } else { "NO" });
+    }
+}
+
+#[allow(dead_code)]
+mod input {
+    use std::collections::VecDeque;
+    use std::io;
+    use std::str::FromStr;
+
+    pub struct Scanner {
+        buffer: VecDeque<String>,
+    }
+
+    impl Scanner {
+        pub fn new() -> Scanner {
+            Scanner {
+                buffer: VecDeque::new(),
+            }
+        }
+
+        pub fn next<T: FromStr>(&mut self) -> T {
+            if self.buffer.is_empty() {
+                let mut input = String::new();
+
+                io::stdin().read_line(&mut input).ok();
+
+                for word in input.split_whitespace() {
+                    self.buffer.push_back(word.to_string())
+                }
+            }
+
+            let front = self.buffer.pop_front().unwrap();
+            front.parse::<T>().ok().unwrap()
+        }
+
+        pub fn next_vec<T: FromStr>(&mut self, n: usize) -> Vec<T> {
+            let mut arr = vec![];
+
+            for _ in 0..n {
+                arr.push(self.next::<T>());
+            }
+
+            arr
+        }
+    }
+}