113 lines
2.7 KiB
Rust
113 lines
2.7 KiB
Rust
use std::rc::Rc;
|
|
use itertools::Itertools;
|
|
|
|
use crate::{ballot::Ballot, header::Header, stats::Stats, util::ScoreItem};
|
|
|
|
#[derive(Debug)]
|
|
pub enum Event {
|
|
Lose(ScoreItem),
|
|
Win(Vec<ScoreItem>),
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct Counter {
|
|
pub header: Rc<Header>,
|
|
pub ballots: Vec<Ballot>,
|
|
pub stats: Stats,
|
|
winners_left: usize,
|
|
enabled: Vec<bool>,
|
|
quota: f64,
|
|
}
|
|
|
|
impl Counter {
|
|
pub fn new<T: std::io::BufRead>(mut csv: quick_csv::Csv<T>, winners: usize) -> Result<Self, Box<dyn std::error::Error>> {
|
|
let header = Rc::new(Header::parse(csv.next().ok_or("csv header missing")??, winners)?);
|
|
let mut ballots = Vec::new();
|
|
let mut stats = Stats::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, &mut stats)?);
|
|
}
|
|
|
|
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,
|
|
stats,
|
|
winners_left,
|
|
enabled,
|
|
quota,
|
|
})
|
|
}
|
|
fn count_primaries(&self) -> Vec<ScoreItem> {
|
|
let mut scores = vec![0.0; self.enabled.len()];
|
|
for ballot in self.ballots.iter() {
|
|
if let Some(index) = ballot.get_primary_index(|id| self.enabled[id]) {
|
|
scores[index] += ballot.get_weight();
|
|
}
|
|
}
|
|
scores.into_iter().enumerate().filter_map(|(index, value)| match self.enabled[index] {
|
|
true => Some(ScoreItem::new(index, value)),
|
|
false => None,
|
|
}).collect_vec()
|
|
}
|
|
pub fn get_quota(&self) -> f64 {
|
|
self.quota
|
|
}
|
|
}
|
|
|
|
impl Iterator for Counter {
|
|
type Item = Event;
|
|
|
|
fn next(&mut self) -> Option<Self::Item> {
|
|
let mut winners = Vec::<ScoreItem>::with_capacity(self.enabled.len());
|
|
let primaries = self.count_primaries();
|
|
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));
|
|
}
|
|
|
|
winners.sort_by(|a,b| b.cmp(a));
|
|
|
|
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))
|
|
}
|
|
}
|
|
|