diff --git a/pyRCV2/method/meek.py b/pyRCV2/method/meek.py index 64a8052..79f859f 100644 --- a/pyRCV2/method/meek.py +++ b/pyRCV2/method/meek.py @@ -39,6 +39,47 @@ class MeekCountCard(CountCard): result.keep_value = keep_value return result +class BallotTree: + def __init__(self): + self.num = Num('0') + self.ballots = [] # List of tuples (ballot, idx) + + self.next_preferences = None # SafeDict: Candidate -> BallotTree + self.next_exhausted = None # BallotTree + + def descend_tree(self): + """Expand one further level of the tree""" + if self.next_exhausted is not None: + raise Exception('Attempt to descend into already descended tree') + + self.next_preferences = SafeDict() + self.next_exhausted = BallotTree() + + for ballot, idx in self.ballots: + if idx is None: + idx = 0 + else: + idx = idx + 1 + + if idx < len(ballot.preferences): + cand = ballot.preferences[idx] + + __pragma__('opov') + if cand not in self.next_preferences: + np = BallotTree() + self.next_preferences[cand] = np + else: + np = self.next_preferences[cand] + + np.num += ballot.value + __pragma__('noopov') + np.ballots.append((ballot, idx)) + else: + __pragma__('opov') + self.next_exhausted.num += ballot.value + __pragma__('noopov') + self.next_exhausted.ballots.append((ballot, idx)) + class MeekSTVCounter(BaseSTVCounter): def describe_options(self): """Overrides BaseSTVCounter.describe_options""" @@ -46,15 +87,23 @@ class MeekSTVCounter(BaseSTVCounter): def __init__(self, *args): BaseSTVCounter.__init__(self, *args) - self.candidates = SafeDict([(c, MeekCountCard()) for c in self.election.candidates]) - # Withdraw candidates + # Convert to MeekCountCard + self.candidates = SafeDict([(c, MeekCountCard()) for c in self.election.candidates]) for candidate in self.election.withdrawn: __pragma__('opov') self.candidates[candidate].state = CandidateState.WITHDRAWN __pragma__('noopov') self._quota_tolerance = Num('1.0001') + + # For tree packing + self.ballots_tree = BallotTree() + for ballot in self.election.ballots: + __pragma__('opov') + self.ballots_tree.num += ballot.value + __pragma__('noopov') + self.ballots_tree.ballots.append((ballot, None)) def reset(self): if self.options['quota_mode'] != 'progressive': @@ -87,6 +136,41 @@ class MeekSTVCounter(BaseSTVCounter): return self.make_result('First preferences') + def distribute_recursively(self, tree, remaining_multiplier): + if tree.next_exhausted is None: + tree.descend_tree() + + # Credit votes at this level + for candidate, cand_tree in tree.next_preferences.items(): + __pragma__('opov') + count_card = self.candidates[candidate] + __pragma__('noopov') + + if count_card.state == CandidateState.HOPEFUL: + # Hopeful candidate has keep value 1, so transfer entire remaining value + __pragma__('opov') + count_card.transfers += remaining_multiplier * cand_tree.num + __pragma__('noopov') + elif count_card.state == CandidateState.EXCLUDED or count_card.state == CandidateState.WITHDRAWN: + # Excluded candidate has keep value 0, so skip over this candidate + # Recurse + self.distribute_recursively(cand_tree, remaining_multiplier) + elif count_card.state == CandidateState.ELECTED: + # Transfer according to elected candidate's keep value + __pragma__('opov') + count_card.transfers += remaining_multiplier * cand_tree.num * count_card.keep_value + new_remaining_multiplier = remaining_multiplier * (Num(1) - count_card.keep_value) + __pragma__('noopov') + # Recurse + self.distribute_recursively(cand_tree, new_remaining_multiplier) + else: + raise STVException('Unexpected candidate state') + + # Credit exhausted votes at this level + __pragma__('opov') + self.exhausted.transfers += remaining_multiplier * tree.next_exhausted.num + __pragma__('noopov') + def distribute_first_preferences(self): """ Overrides BaseSTVCounter.distribute_first_preferences @@ -111,36 +195,7 @@ class MeekSTVCounter(BaseSTVCounter): self.loss_fraction = CountCard() # Distribute votes - for ballot in self.election.ballots: - remaining_value = Num(ballot.value) # Clone the value so we don't update ballot.value - for candidate in ballot.preferences: - __pragma__('opov') - count_card = self.candidates[candidate] - __pragma__('noopov') - - if count_card.state == CandidateState.HOPEFUL: - # Hopeful candidate has keep value 1, so transfer entire remaining value - __pragma__('opov') - count_card.transfers += remaining_value - __pragma__('noopov') - remaining_value = Num(0) - break - elif count_card.state == CandidateState.EXCLUDED or count_card.state == CandidateState.WITHDRAWN: - # Excluded candidate has keep value 0, so skip over this candidate - pass - elif count_card.state == CandidateState.ELECTED: - # Transfer according to elected candidate's keep value - __pragma__('opov') - count_card.transfers += remaining_value * count_card.keep_value - remaining_value *= (Num(1) - count_card.keep_value) - __pragma__('noopov') - else: - raise STVException('Unexpected candidate state') - - # Credit exhausted votes - __pragma__('opov') - self.exhausted.transfers += remaining_value - __pragma__('noopov') + self.distribute_recursively(self.ballots_tree, Num('1')) # Recompute transfers if len(self.step_results) > 0: diff --git a/pyRCV2/safedict/safedict_js.py b/pyRCV2/safedict/safedict_js.py index d8041e7..fe78603 100644 --- a/pyRCV2/safedict/safedict_js.py +++ b/pyRCV2/safedict/safedict_js.py @@ -1,5 +1,5 @@ # pyRCV2: Preferential vote counting -# Copyright © 2020 Lee Yingtong Li (RunasSudo) +# Copyright © 2020–2021 Lee Yingtong Li (RunasSudo) # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Affero General Public License as published by @@ -30,6 +30,9 @@ class SafeDict: def __setitem__(self, key, value): self.impl.set(key, value) + def __contains__(self, key): + return self.impl.has(key) + def items(self): entries = self.impl.entries() # Returns an Iterator return __pragma__('js', 'Array.from(entries)') diff --git a/pyRCV2/safedict/safedict_py.py b/pyRCV2/safedict/safedict_py.py index 1aafc8f..351d746 100644 --- a/pyRCV2/safedict/safedict_py.py +++ b/pyRCV2/safedict/safedict_py.py @@ -29,5 +29,8 @@ class SafeDict: def __setitem__(self, key, value): self.impl[key] = value + def __contains__(self, key): + return key in self.impl + def items(self): return self.impl.items()