Skip to content

Instantly share code, notes, and snippets.

@attgm
Last active August 10, 2023 14:03
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save attgm/6a63c62b2806bd88e1064f862cfced43 to your computer and use it in GitHub Desktop.
Save attgm/6a63c62b2806bd88e1064f862cfced43 to your computer and use it in GitHub Desktop.
Solver for `Independent Set` problem
use proconio::{input, marker::Usize1};
use itertools::Itertools;
const MOD: usize = 10usize.pow(9) + 7;
#[derive(Clone)]
struct Node {
neighbors: Vec<usize>,
value: (usize, usize)
}
impl Node {
fn new() -> Self {
Node {
neighbors: vec![],
value: (1, 1)
}
}
fn add_neighbor(&mut self, i: usize) {
self.neighbors.push(i);
}
fn set_child_value(&mut self, id: usize, value:(usize, usize)) -> bool {
if let Some(index) = self.neighbors.iter().position(|x| *x == id) {
self.neighbors.swap_remove(index);
self.value = ((self.value.0 * value.0) % MOD, (self.value.1 * value.1) % MOD);
}
self.neighbors.len() == 1
}
fn set_parent_node(&mut self) -> Option<usize> {
assert!(self.neighbors.len() <= 1);
self.value = (self.value.1, (self.value.0 + self.value.1) % MOD);
if self.neighbors.len() == 1 {
Some(self.neighbors[0])
} else {
None
}
}
}
fn main(){
input! {
n: usize,
edges: [(Usize1, Usize1); n - 1]
}
let mut nodes = vec![Node::new(); n];
for (i,j) in &edges {
nodes[*i].add_neighbor(*j);
nodes[*j].add_neighbor(*i);
}
let mut queue: Vec<_> = nodes.iter()
.positions(|node| node.neighbors.len() <= 1)
.collect();
let answer = 'main_loop: loop {
let mut next_queue = Vec::new();
for i in &queue {
let index = *i;
if let Some(parent) = nodes[index].set_parent_node() {
let value = nodes[index].value;
if nodes[parent].set_child_value(index, value) {
next_queue.push(parent);
}
} else {
break 'main_loop nodes[index].value.1;
}
}
queue = next_queue;
assert!(queue.len() > 0);
};
println!("{}", answer);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment