diff --git a/src/ballot.rs b/src/ballot.rs index 06ee1a1..30d897e 100644 --- a/src/ballot.rs +++ b/src/ballot.rs @@ -17,6 +17,7 @@ pub struct Ballot { batch_id: usize, paper_id: usize, votes: Vec, + weight: f64, } impl Ballot { @@ -61,12 +62,21 @@ impl Ballot { collection_point_id: collection_point_id.parse()?, batch_id: batch_id.parse()?, paper_id: paper_id.parse()?, + weight: 1.0, votes, }) } - pub fn get_primary_index(&self, enabled: &[bool]) -> Option { + pub fn get_weight(&self) -> f64 { + self.weight + } + pub fn apply_weight(&mut self, weight: f64) { + self.weight *= weight; + } + pub fn get_primary_index(&self, enabled: F) -> Option + where F: Fn(usize) -> bool + { for &id in self.votes.iter() { - if enabled[id] { + if enabled(id) { return Some(id); } } diff --git a/src/counter.rs b/src/counter.rs index 5a6bfdc..f13f579 100644 --- a/src/counter.rs +++ b/src/counter.rs @@ -1,53 +1,102 @@ use std::rc::Rc; +use itertools::Itertools; -use crate::{ballot::Ballot, header::Header}; +use crate::{ballot::Ballot, header::Header, score_item::{ScoreItem, ScoreItems}}; +#[derive(Debug)] +pub enum Event { + Lose(ScoreItem), + Win(Vec), +} #[derive(Debug)] pub struct Counter { pub header: Rc
, pub ballots: Vec, + winners_left: usize, + enabled: Vec, + quota: f64, } impl Counter { - pub fn new(mut csv: quick_csv::Csv) -> Result> { - let header = Header::parse(csv.next().ok_or("csv header missing")??)?; + pub fn new(mut csv: quick_csv::Csv, winners: usize) -> Result> { + let header = Header::parse(csv.next().ok_or("csv header missing")??, winners)?; let mut ballots = Vec::new(); + if winners > header.candidates.len() { + return Err("winners can't be smaller than the candidates list".into()); + } + for row in csv { ballots.push(Ballot::parse(row?, header.clone())?); } + let enabled = vec![true; header.candidates.len()]; + let quota = (ballots.len() as f64) / (header.winners as f64 + 1.0) + 1.0; + let winners_left = header.winners; + Ok(Counter { header, ballots, + winners_left, + enabled, + quota, }) } - - pub fn count_primaries(&self, enabled: &[bool]) -> Vec { - assert!(enabled.len() == self.header.candidates.len()); - - let mut scores = vec![0.0; enabled.len()]; - + fn count_primaries(&self) -> ScoreItems { + let mut scores = vec![0.0; self.enabled.len()]; for ballot in self.ballots.iter() { - if let Some(index) = ballot.get_primary_index(enabled) { - scores[index] += 1.0; + if let Some(index) = ballot.get_primary_index(|id| self.enabled[id]) { + scores[index] += ballot.get_weight(); } } - - scores - } - - pub fn count_all_with_decaying_weights(&self, enabled: &[bool]) -> Vec { - assert!(enabled.len() == self.header.candidates.len()); - - let mut scores = vec![0.0; enabled.len()]; - - for ballot in self.ballots.iter() { - ballot.count_all_with_decaying_weights(&mut scores, enabled); - } - - scores + ScoreItems::new(scores, &self.enabled) } } +impl Iterator for Counter { + type Item = Event; + + fn next(&mut self) -> Option { + let mut winners = Vec::::with_capacity(self.enabled.len()); + let primaries = self.count_primaries().collect_vec(); + let mut lowest = *primaries.first()?; + + for &score in primaries.iter() { + lowest = lowest.min(score); + + if score.value >= self.quota { + winners.push(score); + } + } + + if primaries.len() == 1 && self.winners_left == 1 { + self.enabled[lowest.index] = false; + return Some(Event::Win(vec![lowest])); + } + + if winners.len() == 0 { + self.enabled[lowest.index] = false; + return Some(Event::Lose(lowest)); + } + + for ballot in self.ballots.iter_mut() { + let index = match ballot.get_primary_index(|id| self.enabled[id]) { + Some(v) => v, + None => continue, + }; + if let Some(winner) = winners.iter().filter(|&v| v.index == index).next() { + let surplus = winner.value - self.quota; + ballot.apply_weight(surplus / winner.value); + } + } + + self.winners_left -= winners.len(); + for winner in winners.iter() { + self.enabled[winner.index] = false; + } + + Some(Event::Win(winners)) + } +} + diff --git a/src/header.rs b/src/header.rs index 3980f1b..7d7df40 100644 --- a/src/header.rs +++ b/src/header.rs @@ -9,13 +9,15 @@ use crate::{candidate::{Candidate, CandidateName}, party::Party}; pub struct Header { pub parties: Vec, pub candidates: Vec, + pub winners: usize, } impl Header { - pub fn parse(row: quick_csv::Row) -> Result, Box> { + pub fn parse(row: quick_csv::Row, winners: usize) -> Result, Box> { let mut header = Header { parties: Vec::new(), candidates: Vec::new(), + winners, }; let mut parties_lookup = HashMap::<&str, usize>::new(); diff --git a/src/main.rs b/src/main.rs index ec3f13a..c9e773d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,13 @@ use std::env; use counter::Counter; -use itertools::Itertools; -use util::{percent::PercentDiff, CsvWriter, Percent}; +use util::Percent; pub mod util; pub mod candidate; pub mod ballot; pub mod header; pub mod party; +pub mod score_item; mod counter; fn main() { @@ -20,91 +20,37 @@ fn main() { (args[1].clone(), args[2].parse::().unwrap()) }; let csv = quick_csv::Csv::from_file(csv_path).unwrap().flexible(true); - let counter = Counter::new(csv).unwrap(); + let counter = Counter::new(csv, winner_count).unwrap(); + let mut winners = Vec::with_capacity(winner_count); + let total = counter.ballots.len() as f64; let header = counter.header.clone(); eprintln!("Header: {:#?}", header); eprintln!("Parties: {}", header.parties.len()); eprintln!("Candidates: {}", header.candidates.len()); - let total = counter.ballots.len() as u32; - let mut running = vec![true; header.candidates.len()]; - let mut csv_out = CsvWriter::new(std::io::stdout().lock()); - - { - let mut row = csv_out.row(); - row.skip(4).unwrap(); - for name in header.iter_candidate_name() { - row.write(name).unwrap(); - row.skip(1).unwrap(); - } - } eprintln!("Running election rounds:"); - let mut scores_last: Option> = None; - let mut id_worst_last: Option = None; - - for index in 0..(running.len() - winner_count + 1) { - let scores = counter.count_all_with_decaying_weights(&running); - let scores_ordered = scores.iter().copied().enumerate().filter(|&(id,_)| running[id]).sorted_by(|a, b| f64::total_cmp(&b.1, &a.1)).collect_vec(); - - let (id_worst, score_worst) = match scores_ordered.last() { - Some(&v) => v, - None => break, - }; - - csv_out.row(); - let mut row = csv_out.row(); - - if let Some(scores_last) = scores_last { - let id_worst_last = id_worst_last.unwrap(); - let score_worst_last = scores_last[id_worst_last]; - - row.skip(1).unwrap(); - row.write(header.get_candidate_name(id_worst_last)).unwrap(); - row.write(score_worst_last).unwrap(); - row.write(Percent(score_worst_last, total.into())).unwrap(); - - for ((&score, &score_last), &running) in scores.iter().zip(scores_last.iter()).zip(running.iter()) { - if running { - let score_diff = score - score_last; - row.write(score_diff).unwrap(); - row.write(PercentDiff(score_diff, score_worst_last)).unwrap(); - } else { - row.skip(2).unwrap(); + for (index, event) in (1..).zip(counter) { + match event { + counter::Event::Win(mut v) => { + eprintln!(" {index}: Win:"); + v.sort_by(|a,b| b.cmp(a)); + for winner in v.iter() { + eprintln!(" - {}, {}, {}", header.get_candidate_name(winner.index), winner.value, Percent(winner.value, total)); } + winners.extend_from_slice(&v); } - drop(row); - row = csv_out.row(); - } - - row.write(index + 1).unwrap(); - row.skip(3).unwrap(); - - for (&score, &running) in scores.iter().zip(running.iter()) { - if running { - row.write(score).unwrap(); - row.write(Percent(score, total.into())).unwrap(); - } else { - row.skip(2).unwrap(); + counter::Event::Lose(v) => { + eprintln!(" {index}: Lose: {}. {}, {}", header.get_candidate_name(v.index), v.value, Percent(v.value, total)); } } + } - running[id_worst] = false; - scores_last = Some(scores); - id_worst_last = Some(id_worst); - - let mut total_votes = 0.0; - - eprintln!("Round {}:", index+1); - for (place, &(id, score)) in scores_ordered.iter().take(winner_count).enumerate() { - eprintln!(" {}. {}: {score}, {}", place + 1, header.get_candidate_name(id), Percent(score, total.into())); - total_votes += score; - } - if scores_ordered.len() > winner_count { - eprintln!(" {}. {}: {score_worst}, {}", scores_ordered.len(), header.get_candidate_name(id_worst), Percent(score_worst, total.into())); - } - eprintln!(" Total: {total_votes}, {}", Percent(total_votes.into(), total.into())); + eprintln!(); + eprintln!("Final winners:"); + for (index, winner) in (1..).zip(winners.iter()) { + eprintln!(" {index}: {}", header.get_candidate_name(winner.index)); } } diff --git a/src/score_item.rs b/src/score_item.rs new file mode 100644 index 0000000..4a5388b --- /dev/null +++ b/src/score_item.rs @@ -0,0 +1,61 @@ + +#[derive(Clone,Copy,Debug)] +pub struct ScoreItem { + pub index: usize, + pub value: f64, +} + +pub struct ScoreItems<'a> { + enabled: &'a [bool], + data: Vec, + index: usize, +} + +impl ScoreItem { + pub fn new(index: usize, value: f64) -> Self { + Self { index, value } + } + pub fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.value.total_cmp(&other.value) + } + pub fn max(self, other: Self) -> Self { + if self.value > other.value { + self + } else { + other + } + } + pub fn min(self, other: Self) -> Self { + if self.value > other.value { + other + } else { + self + } + } +} + +impl<'a> ScoreItems<'a> { + pub fn new(data: Vec, enabled: &'a [bool]) -> Self { + Self { data, enabled, index: 0 } + } +} + +impl<'a> Iterator for ScoreItems<'a> { + type Item = ScoreItem; + fn next(&mut self) -> Option { + while self.index < self.data.len() { + if self.enabled[self.index] { + let v = ScoreItem::new(self.index, self.data[self.index]); + self.index += 1; + return Some(v); + } + self.index += 1; + } + None + } + fn size_hint(&self) -> (usize, Option) { + let left = self.data.len() - self.index; + (0, Some(left)) + } +} +