Skip to content

Commit 5855997

Browse files
Add __slots__ and replace PopulatedGraph.degree with direct dictionary lookup
1 parent e12160a commit 5855997

File tree

1 file changed

+16
-9
lines changed

1 file changed

+16
-9
lines changed

gerrychain/tree.py

+16-9
Original file line numberDiff line numberDiff line change
@@ -63,25 +63,32 @@ def uniform_spanning_tree(graph, choice=random.choice):
6363

6464

6565
class PopulatedGraph:
66+
__slots__ = [
67+
"graph",
68+
"subsets",
69+
"population",
70+
"tot_pop",
71+
"ideal_pop",
72+
"epsilon",
73+
"degrees"
74+
]
75+
6676
def __init__(self, graph, populations, ideal_pop, epsilon):
6777
self.graph = graph
6878
self.subsets = {node: {node} for node in graph.node_indicies}
6979
self.population = populations.copy()
7080
self.tot_pop = sum(self.population.values())
7181
self.ideal_pop = ideal_pop
7282
self.epsilon = epsilon
73-
self._degrees = {node: graph.degree(node) for node in graph.node_indicies}
83+
self.degrees = {node: graph.degree(node) for node in graph}
7484

7585
def __iter__(self):
7686
return iter(self.graph)
7787

78-
def degree(self, node):
79-
return self._degrees[node]
80-
8188
def contract_node(self, node, parent):
8289
self.population[parent] += self.population[node]
8390
self.subsets[parent] |= self.subsets[node]
84-
self._degrees[parent] -= 1
91+
self.degrees[parent] -= 1
8592

8693
def has_ideal_population(self, node):
8794
return (
@@ -94,26 +101,26 @@ def has_ideal_population(self, node):
94101

95102
def find_balanced_edge_cuts_contraction(h, choice=random.choice):
96103
# this used to be greater than 2 but failed on small grids:(
97-
root = choice([x for x in h if h.degree(x) > 1])
104+
root = choice([x for x in h if h.degrees[x] > 1])
98105
# BFS predecessors for iteratively contracting leaves
99106
pred = predecessors(h.graph, root)
100107

101108
cuts = []
102-
leaves = deque(x for x in h if h.degree(x) == 1)
109+
leaves = deque(x for x in h if h.degrees[x] == 1)
103110
while len(leaves) > 0:
104111
leaf = leaves.popleft()
105112
if h.has_ideal_population(leaf):
106113
cuts.append(Cut(edge=(leaf, pred[leaf]), subset=h.subsets[leaf].copy()))
107114
# Contract the leaf:
108115
parent = pred[leaf]
109116
h.contract_node(leaf, parent)
110-
if h.degree(parent) == 1 and parent != root:
117+
if h.degrees[parent] == 1 and parent != root:
111118
leaves.append(parent)
112119
return cuts
113120

114121

115122
def find_balanced_edge_cuts_memoization(h, choice=random.choice):
116-
root = choice([x for x in h if h.degree(x) > 1])
123+
root = choice([x for x in h if h.degrees[x] > 1])
117124
pred = predecessors(h.graph, root)
118125
succ = successors(h.graph, root)
119126
total_pop = h.tot_pop

0 commit comments

Comments
 (0)