Use tree-packed ballots for Meek STV

Runtime on Senate data:
1m54s not tree packed
13s tree packed (!!!!!)

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:
+		self.distribute_recursively(self.ballots_tree, Num('1'))
-			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')
 		
 		# Recompute transfers
 		if len(self.step_results) > 0:

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)')

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()