Cod sursa(job #2692587)

Utilizator CyborgSquirrelJardan Andrei CyborgSquirrel Data 3 ianuarie 2021 11:43:45
Problema Arbore partial de cost minim Scor 0
Compilator rs Status done
Runda Arhiva educationala Marime 2.04 kb
use std::io::BufReader;
use std::io::prelude::*;
use std::fs::File;
use std::io::Write;
use std::fs::OpenOptions;

#[derive(Copy, Clone, Debug)]
struct Edge{
	a:usize,
	b:usize,
	cost:i32,
}

struct Tree{
	v:Vec::<usize>,
}
impl Tree{
	fn dad(&mut self, a: usize) -> usize{
		if self.v[a] != a {
			self.v[a] = self.dad(self.v[a]);
		}
		self.v[a]
	}
	fn unite(&mut self, a: usize, b: usize){
		let a = self.dad(a);
		let b = self.dad(b);
		self.v[a] = b;
	}
	fn check(&mut self, a: usize, b: usize) -> bool{
		let a = self.dad(a);
		let b = self.dad(b);
		a == b
	}
}

fn main() -> std::io::Result<()>{
	let file = File::open("apm.in")?;
	let file = BufReader::new(file);
	
	let mut lines = file
		.lines();
	
	let line = lines
		.next()
		.unwrap().unwrap();
	let mut line = line
		.split_whitespace();
	
	let n = line.next().unwrap().parse::<i32>().unwrap();
	let m = line.next().unwrap().parse::<i32>().unwrap();
	
	let mut edges = Vec::<Edge>::with_capacity(400041);
	let mut tree = Tree{
		v: {
			let mut v = Vec::<usize>::with_capacity(200041);
			for i in 0..=n{
				v.push(i as usize);
			}
			v
		}
	};
	
	let mut used_edges = Vec::<&Edge>::with_capacity(200041);
	let mut cost: i32 = 0;
	
	for _i in 0..m{
		let line = lines
			.next()
			.unwrap().unwrap();
		let mut line = line
			.split_whitespace();
		
		edges.push(
			Edge{
				a: line.next().unwrap().parse::<usize>().unwrap(),
				b: line.next().unwrap().parse::<usize>().unwrap(),
				cost: line.next().unwrap().parse::<i32>().unwrap(),
			}
		);
	}
	
	edges.sort_by(|a, b| a.cost.partial_cmp(&b.cost).unwrap());
	
	for edge in edges.iter() {
		if !tree.check(edge.a, edge.b) {
			tree.unite(edge.a, edge.b);
			used_edges.push(&edge);
			cost += edge.cost;
		}
	}
	
	let mut file = OpenOptions::new()
		.truncate(true)
		.write(true)
		.open("apm.out")?;
	
	write!(file, "{}\n", cost)?;
	write!(file, "{}\n", used_edges.len())?;
	for edge in used_edges {
		write!(file, "{} {}\n", edge.a, edge.b)?;
	}
	
    Ok(())
}