def county_bipartition_tree(graph,
pop_col,
pop_target,
epsilon,
county_col="COUNTYFP10",
node_repeats=1,
spanning_tree=None,
choice=random.choice,
attempts_before_giveup=100):
populations = {node: graph.nodes[node][pop_col] for node in graph}
balanced_subtree = None
if spanning_tree is None:
spanning_tree = county_random_spanning_tree(graph,
county_col=county_col)
restarts = 0
counter = 0
while balanced_subtree is None and counter < attempts_before_giveup:
# print(counter)
if restarts == node_repeats:
spanning_tree = county_random_spanning_tree(graph,
county_col=county_col)
restarts = 0
counter += 1
h = PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
if first_check_counties and restarts == 0:
balanced_subtree = split_tree_at_county(h,
choice=choice,
county_col=county_col)
if balanced_subtree is None:
h = PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
balanced_subtree = contract_leaves_until_balanced_or_none(
h, choice=choice)
restarts += 1
if counter >= attempts_before_giveup:
return set()
return balanced_subtree
def recursive_my_mst_k_partition_tree_random(
test_graph,
pop_col,
pop_target,
epsilon,
num_blocks,
our_blocks,
spanning_tree=None,
choice=random.choice,):
populations = {node: test_graph.nodes[node][pop_col] for node in test_graph}
# pop_target = np.sum([test_graph.nodes[node]["population"] for node in test_graph]) / num_blocks
possible_cuts = []
while len(possible_cuts) == 0:
spanning_tree = get_spanning_tree_mst(test_graph)
h = PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
possible_cuts = find_balanced_edge_cuts(h, choice=choice)
# remove possible cuts to get new partition
# spanning_tree.remove_edge(possible_cuts[0].edge[0], possible_cuts[0].edge[1])
# unfrozen_tree = nx.Graph(spanning_tree)
# comps = nx.connected_components(unfrozen_tree)
#
# list_comps = list(comps)
# population_block = 0
# list_blocks = []
#
# # iterate through all new blocks to find out which block's population close to ideal population, add to block list
# for x in list_comps:
# for y in x:
# population_block += populations[y]
# list_blocks.append(population_block)
# population_block = 0
# index = min(range(len(list_blocks)), key=lambda x: abs(list_blocks[x]-pop_target))
# our_blocks.append(list_comps[index])
# sub_graph = test_graph.copy()
removed_set = random.choice(possible_cuts).subset
our_blocks.append(removed_set)
sub_graph = test_graph.copy()
# for x in list_comps[index]:
# sub_graph.remove_node(x)
for x in removed_set:
sub_graph.remove_node(x)
return sub_graph
def test_find_balanced_cuts_contraction():
tree = networkx.Graph([(0, 1), (1, 2), (1, 4), (3, 4), (4, 5), (3, 6),
(6, 7), (6, 8)])
# 0 - 1 - 2
# ||
# 3= 4 - 5
# ||
# 6- 7
# |
# 8
populated_tree = PopulatedGraph(tree, {node: 1
for node in tree},
len(tree) / 2, 0.5)
cuts = find_balanced_edge_cuts_contraction(populated_tree)
edges = set(tuple(sorted(cut.edge)) for cut in cuts)
assert edges == {(1, 4), (3, 4), (3, 6)}
def my_uu_bipartition_tree_random(graph,
pop_col,
pop_target,
epsilon,
node_repeats=1,
spanning_tree=None,
choice=random.choice):
populations = {node: graph.nodes[node][pop_col] for node in graph}
possible_cuts = []
if spanning_tree is None:
spanning_tree = get_spanning_tree_u_w(graph)
while len(possible_cuts) == 0:
spanning_tree = get_spanning_tree_u_w(graph)
h = PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
possible_cuts = find_balanced_edge_cuts(h, choice=choice)
return choice(possible_cuts).subset
def recursive_my_mst_k_partition_tree_random(
test_graph,
pop_col,
pop_target,
epsilon,
num_blocks,
our_blocks,
spanning_tree=None,
choice=random.choice,):
populations = {node: test_graph.nodes[node][pop_col] for node in test_graph}
pop_target = np.sum([test_graph.nodes[node]["population"] for node in test_graph]) / num_blocks
possible_cuts = []
while len(possible_cuts) == 0:
spanning_tree = get_spanning_tree_mst(test_graph)
h = PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
possible_cuts = find_balanced_edge_cuts(h, choice=choice)
spanning_tree.remove_edge(possible_cuts[0].edge[0], possible_cuts[0].edge[1])
unfrozen_tree = nx.Graph(spanning_tree)
comps = nx.connected_components(unfrozen_tree)
list_comps = list(comps)
population_block = 0
list_blocks = []
for x in list_comps:
for y in x:
population_block += populations[y]
list_blocks.append(population_block)
population_block = 0
index = min(range(len(list_blocks)), key=lambda x: abs(list_blocks[x]-pop_target))
our_blocks.append(list_comps[index])
sub_graph = test_graph.copy()
for x in list_comps[index]:
sub_graph.remove_node(x)
return sub_graph