#586942 (Rust) No.1301 Strange Graph Shortest Path

提出ソース
結果

問題	No.1301 Strange Graph Shortest Path
コンテスト
ユーザー	nebocco
提出日時	2020-11-29 11:52:15
言語	Rust (1.83.0 + proconio)
結果	TLE
実行時間	-
コード長	17,921 bytes
コンパイル時間	15,303 ms
コンパイル使用メモリ	399,700 KB
実行使用メモリ	57,060 KB
最終ジャッジ日時	2024-09-13 01:57:00
合計ジャッジ時間	19,409 ms
ジャッジサーバーID （参考情報）	judge3 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 2
other	TLE * 1 -- * 32
権限があれば一括ダウンロードができます
ソースコード

raw source code

fn main() {
    let mut io = IO::new();
    let (n, m): (usize, usize) = io.scan();
    let mut ns: NetworkSimplex<i32, i64> = NetworkSimplex::new();
    for _ in 0..m {
        let (u, v, c, d): (usize, usize, i64, i64) = io.scan();
        let u = u - 1;
        let v = v - 1;
        ns.add_edge(u, v, 0, 1, c);
        ns.add_edge(v, u, 0, 1, c);
        ns.add_edge(u, v, 0, 1, d);
        ns.add_edge(v, u, 0, 1, d);
    }
    ns.add_supply(0, 2);
    ns.add_demand(n-1, 2);
    let ret = ns.run().unwrap();
    let ans = ret.get_value::<i64>();
    io.println(ans);
}



// ------------ traits start ------------

use std::fmt::Display;

pub trait Cost:
	Element
	+ Display
    + Copy
    + Eq
    + Ord
    + Zero
    + One
    + Add<Output = Self>
    + AddAssign
    + Sub<Output = Self>
    + Neg<Output = Self>
{
    fn is_positive(&self) -> bool {
        self > &Self::zero()
    }
    fn is_negative(&self) -> bool {
        self < &Self::zero()
    }
}

pub trait Flow: Cost + SubAssign {
    fn abs(&self) -> Self {
        if self.is_negative() {
            -*self
        } else {
            *self
        }
    }
}

macro_rules! impl_flow {
    ($($T:ty,)*) => {
		$(
			impl Flow for $T {}
			impl Cost for $T {}
		)*
    };
}

impl_flow!(
	i8, i16, i32, i64, i128, isize,
);

// ------------ traits end ------------

use std::cmp::{ max, min };
use std::collections::HashSet;

struct Edge<F, C> {
    src: usize,
    dst: usize,
    flow: F,
    capacity: F,
    cost: C,
}

#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Debug, Hash)]
pub struct EdgeId(usize);
impl EdgeId {
    fn rev(self) -> Self {
        EdgeId(self.0 ^ 1)
    }
}

struct VertexData<C> {
    potential: C,
    adjacent_edges: Vec<EdgeId>,
    parent: Option<usize>,
    parent_edge: Option<EdgeId>, // out-tree, i.e. this node == e.src
    depth: usize,
    tree_edges: HashSet<EdgeId>,
}
impl<C: Zero> Default for VertexData<C> {
    fn default() -> Self {
        Self {
            potential: C::zero(),
            adjacent_edges: Vec::new(),
            parent: None,
            parent_edge: None,
            depth: 0,
            tree_edges: HashSet::new(),
        }
    }
}
#[derive(Default)]
pub struct NetworkSimplex<F: Flow, C: Cost> {
    edges: Vec<Edge<F, C>>,
    balances: Vec<F>,
}
struct TemporaryData<C: Cost> {
    vertices: Vec<VertexData<C>>,
    n: usize,
    root: usize,
    block_size: usize,
    next_scan_start: usize,
}

pub struct Ret<F, C> {
    edges: Vec<(F, C)>,
    potential: Vec<C>,
}
impl<F: Flow, C: Cost> Ret<F, C> {
    pub fn get_value<T>(&self) -> T
    where
        T: From<F> + From<C> + Mul<Output = T> + Add<Output = T> + Zero,
    {
        self.edges
            .iter()
            .filter(|(f, _)| f.is_positive())
            .map(|(f, c)| T::from(*f) * T::from(*c))
            .fold(T::zero(), |a, b| a + b)
    }
    pub fn get_flow(&self, e: EdgeId) -> F {
        self.edges[e.0].0
    }
    pub fn get_potential(&self, v: usize) -> C {
        self.potential[v]
    }
}

impl<F: Flow, C: Cost> NetworkSimplex<F, C> {
    pub fn new() -> Self {
        Self {
            edges: Vec::new(),
            balances: Vec::new(),
        }
    }

    pub fn add_edge(&mut self, src: usize, dst: usize, lower: F, upper: F, cost: C) -> EdgeId {
        assert!(
            lower <= upper,
            "lower {} should be less or equal to upper {}",
            lower,
            upper
        );
        let id = self.edges.len();
        self.edges.push(Edge {
            src,
            dst,
            flow: lower,
            capacity: upper,
            cost,
        });
        self.edges.push(Edge {
            src: dst,
            dst: src,
            flow: -lower,
            capacity: -lower,
            cost: -cost,
        });
        if !lower.is_zero() {
            self.add_demand(src, lower);
            self.add_supply(dst, lower);
        }
        EdgeId(id)
    }

    pub fn add_supply(&mut self, v: usize, b: F) {
        let n = max(v + 1, self.balances.len());
        self.balances.resize_with(n, F::zero);
        self.balances[v] += b;
    }

    pub fn add_demand(&mut self, v: usize, b: F) {
        self.add_supply(v, -b);
    }

    fn get_edge(&self, e: EdgeId) -> &Edge<F, C> {
        &self.edges[e.0]
    }

    fn get_edge_mut(&mut self, e: EdgeId) -> &mut Edge<F, C> {
        &mut self.edges[e.0]
    }

    /// return true iff this was a saturating push
    fn add_flow(&mut self, e: EdgeId, f: F) -> bool {
        self.get_edge_mut(e.rev()).flow -= f;
        let e = self.get_edge_mut(e);
        e.flow += f;
        e.flow == e.capacity
    }

    fn residual_capacity(e: &Edge<F, C>) -> F {
        e.capacity - e.flow
    }

    fn reduced_cost(data: &TemporaryData<C>, e: &Edge<F, C>) -> C {
        e.cost + data.vertices[e.src].potential - data.vertices[e.dst].potential
    }

    fn update_tree(&self, data: &mut TemporaryData<C>, v: usize) {
        let mut stack = vec![v];
        while let Some(v) = stack.pop() {
            let adj = std::mem::take(&mut data.vertices[v].tree_edges);
            for &eid in adj.iter() {
                let e = self.get_edge(eid);
                if data.vertices[v].parent == Some(e.dst) {
                    continue;
                }
                data.vertices[e.dst].parent = Some(v);
                data.vertices[e.dst].parent_edge = Some(eid.rev());
                data.vertices[e.dst].depth = data.vertices[e.src].depth + 1;
                data.vertices[e.dst].potential = data.vertices[e.src].potential + e.cost;
                stack.push(e.dst);
            }
            data.vertices[v].tree_edges = adj;
        }
    }

    fn prepare_data(&mut self) -> TemporaryData<C> {
        // allocate root vertex
        let mut infinity = C::one();
        let mut data = TemporaryData {
            vertices: Default::default(),
            n: self.balances.len(),
            root: 0,
            block_size: 1,
            next_scan_start: 0,
        };

        data.vertices.clear();
        for (i, e) in self.edges.iter().enumerate() {
            data.n = max(data.n, 1 + e.src);
            data.vertices.resize_with(data.n, Default::default);
            data.vertices[e.src].adjacent_edges.push(EdgeId(i));
            if e.cost.is_positive() {
                infinity += e.cost;
            }
        }
        data.root = data.n;
        data.n += 1;
        let root = data.root;
        data.vertices.resize_with(data.n, Default::default);
        self.balances.resize_with(data.n - 1, F::zero);
        for v in 0..root {
            let b = std::mem::replace(&mut self.balances[v], F::zero());
            let (x, y, cap) = if b.is_negative() {
                (root, v, -b)
            } else {
                (v, root, b + F::one())
            };
            let eid = self.add_edge(x, y, F::zero(), cap, infinity);
            self.add_flow(eid, b.abs());
            data.vertices[x].adjacent_edges.push(eid);
            data.vertices[y].adjacent_edges.push(eid.rev());
            data.vertices[x].tree_edges.insert(eid);
            data.vertices[y].tree_edges.insert(eid.rev());
        }
        data.block_size = min(
            (self.edges.len() as f64).sqrt() as usize + 10,
            self.edges.len(),
        );
        self.update_tree(&mut data, root);
        data
    }

    fn select_edge(&mut self, data: &mut TemporaryData<C>) -> Option<EdgeId> {
        let mut edges = (data.next_scan_start..self.edges.len())
            .chain(0..data.next_scan_start)
            .map(EdgeId)
            .peekable();
        while edges.peek().is_some() {
            let mut selection = Option::None;
            for _ in 0..data.block_size {
                match edges.next() {
                    None => {
                        break;
                    }
                    Some(id) => {
                        let e = self.get_edge_mut(id);
                        if e.flow == e.capacity {
                            continue;
                        }
                        let rc = Self::reduced_cost(data, e);
                        if rc.is_negative() {
                            let candidate = (rc, id);
                            if let Some(current) = selection.take() {
                                selection = Some(min(current, candidate))
                            } else {
                                selection = Some(candidate)
                            }
                        }
                    }
                }
            }
            if let Some((_, eid)) = selection {
                if let Some(nid) = edges.peek() {
                    data.next_scan_start = nid.0;
                }
                return Some(eid);
            }
        }
        None
    }

    fn pivot(&mut self, data: &mut TemporaryData<C>, eid: EdgeId) {
        let entering_edge = self.get_edge(eid);
        let Edge { src, dst, .. } = *entering_edge;
        let mut f = Self::residual_capacity(entering_edge);
        let mut a = src;
        let mut b = dst;
        while a != b {
            if data.vertices[a].depth > data.vertices[b].depth {
                let down_edge = data.vertices[a].parent_edge.unwrap().rev();
                let e = self.get_edge(down_edge);
                f = min(f, Self::residual_capacity(e));
                a = e.src;
            } else {
                let up_edge = data.vertices[b].parent_edge.unwrap();
                let e = self.get_edge(up_edge);
                f = min(f, Self::residual_capacity(e));
                b = e.dst;
            }
        }
        enum LeavingSide {
            SRC,
            DST,
            ENTER,
        }
        let mut leaving_side = LeavingSide::ENTER;
        let top = a;
        let mut leaving_edge_id = None;
        a = src;
        while a != top {
            let v_data = &data.vertices[a];
            let down_edge = v_data.parent_edge.unwrap().rev();
            if self.add_flow(down_edge, f) && leaving_edge_id.is_none() {
                leaving_edge_id = Some(down_edge);
                leaving_side = LeavingSide::SRC;
            }
            a = v_data.parent.unwrap();
        }
        if self.add_flow(eid, f) {
            leaving_edge_id = Some(eid);
            leaving_side = LeavingSide::ENTER;
        }
        b = dst;
        while b != top {
            let v_data = &data.vertices[b];
            let up_edge = v_data.parent_edge.unwrap();
            if self.add_flow(up_edge, f) {
                leaving_edge_id = Some(up_edge);
                leaving_side = LeavingSide::DST;
            }
            b = v_data.parent.unwrap();
        }
        let leaving_edge_id = leaving_edge_id.unwrap();
        let leaving_e = self.get_edge(leaving_edge_id);
        if leaving_edge_id == eid {
            return;
        }
        assert!(data.vertices[src].tree_edges.insert(eid));
        assert!(data.vertices[dst].tree_edges.insert(eid.rev()));
        assert!(data.vertices[leaving_e.src]
            .tree_edges
            .remove(&leaving_edge_id));
        assert!(data.vertices[leaving_e.dst]
            .tree_edges
            .remove(&leaving_edge_id.rev()));
        match leaving_side {
            LeavingSide::SRC => self.update_tree(data, dst),
            LeavingSide::DST => self.update_tree(data, src),
            LeavingSide::ENTER => (),
        }
    }

    pub fn run(&mut self) -> Option<Ret<F, C>> {
        let mut data = self.prepare_data();
        while let Some(eid) = self.select_edge(&mut data) {
            self.pivot(&mut data, eid);
        }
        for e in self.edges.split_off(self.edges.len() - 2 * (data.n - 1)) {
            if !e.flow.is_zero() {
                return None;
            }
        }
        Some(Ret {
            edges: self.edges.iter().map(|e| (e.flow, e.cost)).collect(),
            potential: data
                .vertices
                .iter()
                .take(data.n - 1)
                .map(|v| v.potential)
                .collect(),
        })
    }
}


// ------------ libraries start ------------

// ------------ libraries end ------------

// ------------ traits start ------------

use std::marker::Sized;
use std::ops::*;

/// 元
pub trait Element: Sized + Clone + PartialEq {}
impl<T: Sized + Clone + PartialEq> Element for T {}

/// 結合性
pub trait Associative: Magma {}

/// マグマ
pub trait Magma: Element + Add<Output=Self> {}
impl<T: Element + Add<Output=Self>> Magma for T {}

/// 半群
pub trait SemiGroup: Magma + Associative {}
impl<T: Magma + Associative> SemiGroup for T {}

/// モノイド
pub trait Monoid: SemiGroup + Zero {}
impl<T: SemiGroup + Zero> Monoid for T {}

pub trait ComMonoid: Monoid + AddAssign {}
impl<T: Monoid + AddAssign> ComMonoid for T {}

/// 群
pub trait Group: Monoid + Neg<Output=Self> {}
impl<T: Monoid + Neg<Output=Self>> Group for T {}

pub trait ComGroup: Group + ComMonoid {}
impl<T: Group + ComMonoid> ComGroup for T {}

/// 半環
pub trait SemiRing: ComMonoid + Mul<Output=Self> + One {}
impl<T: ComMonoid + Mul<Output=Self> + One> SemiRing for T {}

/// 環
pub trait Ring: ComGroup + SemiRing {}
impl<T: ComGroup + SemiRing> Ring for T {}

pub trait ComRing: Ring + MulAssign {}
impl<T: Ring + MulAssign> ComRing for T {}

/// 体
pub trait Field: ComRing + Div<Output=Self> + DivAssign {}
impl<T: ComRing + Div<Output=Self> + DivAssign> Field for T {}

/// 加法単元
pub trait Zero: Element {
    fn zero() -> Self;
    fn is_zero(&self) -> bool;
}

/// 乗法単元
pub trait One: Element {
    fn one() -> Self;
    fn is_one(&self) -> bool;
}

macro_rules! impl_integer {
    ($($T:ty,)*) => {
        $(
            impl Associative for $T {}

            impl Zero for $T {
                fn zero() -> Self { 0 }
                fn is_zero(&self) -> bool { *self == 0 }
            }

            impl<'a> Zero for &'a $T {
                fn zero() -> Self { &0 }
                fn is_zero(&self) -> bool { *self == &0 }
            }

            impl One for $T {
                fn one() -> Self { 1 }
                fn is_one(&self) -> bool { *self == 1 }
            }
            impl<'a> One for &'a $T {
                fn one() -> Self { &1 }
                fn is_one(&self) -> bool { *self == &1 }
            }
        )*
    };
}
impl_integer! {
    i8, i16, i32, i64, i128, isize,
    u8, u16, u32, u64, u128, usize,
}

// ------------ traits end ------------

// ------------ io module start ------------

use std::io::{stdout, BufWriter, Read, StdoutLock, Write};

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

impl IO {
	pub fn new() -> Self {
		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()),
		}
	}
	fn scan_str(&mut self) -> &'static str {
		self.iter.next().unwrap()
	}
	fn scan_raw(&mut self) -> &'static [u8] {
		self.scan_str().as_bytes()
	}
	pub fn scan<T: Scan>(&mut self) -> T {
		T::scan(self)
	}
	pub fn scan_vec<T: Scan>(&mut self, n: usize) -> Vec<T> {
		(0..n).map(|_| self.scan()).collect()
	}
}

impl IO {
	pub fn print<T: Print>(&mut self, x: T) {
		T::print(self, x);
	}
	pub fn println<T: Print>(&mut self, x: T) {
		self.print(x);
		self.print("\n");
	}
	pub fn iterln<T: Print, I: Iterator<Item = T>>(&mut self, mut iter: I, delim: &str) {
		if let Some(v) = iter.next() {
			self.print(v);
			for v in iter {
				self.print(delim);
				self.print(v);
			}
		}
		self.print("\n");
	}
	pub fn flush(&mut self) {
		self.buf.flush().unwrap();
	}
}

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

pub trait Scan {
	fn scan(io: &mut IO) -> Self;
}

macro_rules! impl_parse_int {
	($($t:tt),*) => {
		$(
			impl Scan for $t {
				fn scan(s: &mut IO) -> Self {
					let mut res = 0;
					let mut neg = false;
					for d in s.scan_raw() {
						if *d == b'-' {
							neg = true;
						} else {
							res *= 10;
							res += (*d - b'0') as $t;
						}
					}
					if neg { res = res.wrapping_neg(); }
					res
				}
			}
		)*
	};
}

impl_parse_int!(i16, i32, i64, isize, u16, u32, u64, usize);

impl<T: Scan, U: Scan> Scan for (T, U) {
	fn scan(s: &mut IO) -> Self {
		(T::scan(s), U::scan(s))
	}
}

impl<T: Scan, U: Scan, V: Scan> Scan for (T, U, V) {
	fn scan(s: &mut IO) -> Self {
		(T::scan(s), U::scan(s), V::scan(s))
	}
}

impl<T: Scan, U: Scan, V: Scan, W: Scan> Scan for (T, U, V, W) {
	fn scan(s: &mut IO) -> Self {
		(T::scan(s), U::scan(s), V::scan(s), W::scan(s))
	}
}

pub trait Print {
	fn print(w: &mut IO, x: Self);
}

macro_rules! impl_print_int {
	($($t:ty),*) => {
		$(
			impl Print for $t {
				fn print(w: &mut IO, x: Self) {
					w.buf.write_all(x.to_string().as_bytes()).unwrap();
				}
			}
		)*
	};
}

impl_print_int!(i16, i32, i64, isize, u16, u32, u64, usize);

impl Print for u8 {
	fn print(w: &mut IO, x: Self) {
		w.buf.write_all(&[x]).unwrap();
	}
}

impl Print for &[u8] {
	fn print(w: &mut IO, x: Self) {
		w.buf.write_all(x).unwrap();
	}
}

impl Print for &str {
	fn print(w: &mut IO, x: Self) {
		w.print(x.as_bytes());
	}
}

impl<T: Print, U: Print> Print for (T, U) {
	fn print(w: &mut IO, (x, y): Self) {
		w.print(x);
		w.print(" ");
		w.print(y);
	}
}

impl<T: Print, U: Print, V: Print> Print for (T, U, V) {
	fn print(w: &mut IO, (x, y, z): Self) {
		w.print(x);
		w.print(" ");
		w.print(y);
		w.print(" ");
		w.print(z);
	}
}

impl<T: Print, U: Print, V: Print, W: Print> Print for (T, U, V, W) {
	fn print(w: &mut IO, (x, y, z, a): Self) {
		w.print(x);
		w.print(" ");
		w.print(y);
		w.print(" ");
        w.print(z);
        w.print(" ");
		w.print(a);
	}
}

// ------------ io module end ------------
yukicoder

結果

ソースコード
