use hashbrown::HashMap;
use num_traits::cast::NumCast;
use num_traits::Num;
use std::hash::Hash;
use std::ops::AddAssign;
use super::result::CountedCandidates;
use super::result::RankedWinners;
pub type DefaultPluralityTally<T> = PluralityTally<T, u64>;
pub struct PluralityTally<T, C = u64>
where
T: Eq + Clone + Hash, C: Copy + PartialOrd + AddAssign + Num + NumCast, {
running_total: HashMap<T, C>,
num_winners: u32,
}
impl<T, C> PluralityTally<T, C>
where
T: Eq + Clone + Hash, C: Copy + PartialOrd + AddAssign + Num + NumCast, {
pub fn new(num_winners: u32) -> Self {
return PluralityTally {
running_total: HashMap::new(),
num_winners: num_winners,
};
}
pub fn with_capacity(num_winners: u32, expected_candidates: usize) -> Self {
return PluralityTally {
running_total: HashMap::with_capacity(expected_candidates),
num_winners: num_winners,
};
}
pub fn add(&mut self, vote: T) {
self.add_weighted(vote, C::one());
}
pub fn add_ref(&mut self, vote: &T) {
self.add_weighted_ref(vote, C::one());
}
pub fn add_weighted(&mut self, vote: T, weight: C) {
*self.running_total.entry(vote).or_insert(C::zero()) += weight;
}
pub fn add_weighted_ref(&mut self, vote: &T, weight: C) {
if self.running_total.contains_key(&vote) {
if let Some(x) = self.running_total.get_mut(&vote) {
*x += weight;
}
} else {
self.running_total.insert(vote.clone(), weight);
}
}
pub fn candidates(&self) -> Vec<T> {
return self.running_total.keys().map(|k| k.clone()).collect();
}
pub fn winners(&self) -> RankedWinners<T> {
return self.get_counted().into_ranked(self.num_winners);
}
pub fn totals(&self) -> Vec<(T, C)> {
return self.get_counted().into_vec();
}
pub fn ranked(&self) -> Vec<(T, u32)> {
return self.get_counted().into_ranked(0).into_vec();
}
fn get_counted(&self) -> CountedCandidates<T, C> {
let mut counted = CountedCandidates::new();
for (candidate, votecount) in self.running_total.iter() {
counted.push(candidate.clone(), *votecount);
}
return counted;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn plurality_test() {
let mut tally = DefaultPluralityTally::new(2);
tally.add("Alice");
tally.add("Cir");
tally.add("Bob");
tally.add("Alice");
tally.add("Alice");
tally.add("Bob");
assert_eq!(tally.candidates().len(), 3);
assert_eq!(tally.totals(), vec![("Alice", 3), ("Bob", 2), ("Cir", 1)]);
assert_eq!(tally.ranked(), vec![("Alice", 0), ("Bob", 1), ("Cir", 2)]);
let winners = tally.winners();
assert_eq!(winners.contains(&"Alice"), true);
assert_eq!(winners.contains(&"Bob"), true);
assert_eq!(winners.contains(&"Cir"), false);
assert_eq!(winners.contains(&"Rando"), false);
let mut tally = DefaultPluralityTally::new(1);
tally.add(99);
tally.add(100);
tally.add(99);
tally.add(99);
tally.add(1);
tally.add(1);
tally.add(2);
tally.add(0);
let winners = tally.winners();
assert_eq!(winners.contains(&99), true);
assert_eq!(winners.contains(&100), false);
assert_eq!(winners.contains(&1), false);
assert_eq!(winners.contains(&2), false);
assert_eq!(winners.contains(&1000), false);
let mut tally = DefaultPluralityTally::with_capacity(1, 2);
let candidate_id_1 = 123;
let candidate_id_2 = 456;
tally.add_ref(&candidate_id_1);
tally.add_ref(&candidate_id_2);
let winners = tally.winners();
assert_eq!(winners.contains(&candidate_id_1), true);
assert_eq!(winners.contains(&candidate_id_2), true);
}
}