diff --git a/cidr4_merger.py b/cidr4_merger.py index a25982a..b1fd20f 100644 --- a/cidr4_merger.py +++ b/cidr4_merger.py @@ -1,4 +1,5 @@ import cProfile +from collections import defaultdict Node = tuple[int, int, int, int] @@ -113,16 +114,71 @@ def make_cidr4(ip, mask_len) -> str: return f"{ip_address}/{mask_len}" -def answer(nodes: list[Node], required_len: int) -> tuple[list[str], int]: - nodes = sort_nodes(nodes) - merged_nodes = merge_nodes(nodes, required_len) +def lift_lonely_node(nodes: list[Node], singles: list[Node]) -> list[Node]: + # find single whose parent has the least added addresses + min_single, min_parent = singles[0], make_parent(singles[0]) + for node in singles[1:]: + parent = make_parent(node) + if parent[2] < min_parent[2]: + min_single, min_parent = node, parent - cidr4s = [] - sum_added_ips = 0 - for ip_value, mask_len, added_ips, _ in merged_nodes: - cidr4s.append(make_cidr4(ip_value, mask_len)) - sum_added_ips += added_ips - return cidr4s, sum_added_ips + new_nodes = [x for x in nodes] + new_nodes.remove(min_single) + new_nodes.append(min_parent) + new_nodes = sort_nodes(new_nodes) + return new_nodes + + +def merge_neighbors( + nodes: list[Node], neighbours: list[tuple[Node, Node]] +) -> list[Node]: + new_nodes = [x for x in nodes] + for a, b in neighbours: + parent = make_parent(a, b) + new_nodes.remove(a) + new_nodes.remove(b) + new_nodes.append(parent) + return sort_nodes(nodes) + + +def find_neighbours_singles(groups: defaultdict) -> tuple[list, list]: + neighbours = [] + singles = [] + for group in groups.values(): + i = 0 + while i < len(group) - 1: + a, b = group[i], group[i + 1] + ip_a, mask_len_a, _, parent_ip_a = a + ip_b, mask_len_b, _, parent_ip_b = b + if have_same_parent(mask_len_a, parent_ip_a, mask_len_b, parent_ip_b): + neighbours.append((a, b)) + i += 2 + else: + singles.append(a) + i += 1 + if i == len(group) - 1: + singles.append(group[i]) + return neighbours, singles + + +def make_groups(nodes: list[Node]) -> defaultdict: + groups = defaultdict(list) + for n in nodes: + groups[n[1]].append(n) + return groups + + +def merge_nodes_recursion(nodes: list[Node], required_len: int) -> list[Node]: + if len(nodes) <= required_len: + return nodes + groups = make_groups(nodes) + neighbours, singles = find_neighbours_singles(groups) + print(f"{len(nodes)=} {len(singles)=} {len(neighbours)=}") + if neighbours: + new_nodes = merge_neighbors(nodes, neighbours) + return merge_nodes_recursion(new_nodes, required_len) + new_nodes = lift_lonely_node(nodes, singles) + return merge_nodes_recursion(new_nodes, required_len) def main(): @@ -131,7 +187,16 @@ def main(): data = get_data(file) nodes = list(map(cidr4_to_node, data)) - cidr4s, sum_added_ips = answer(nodes, required_len) + + nodes = sort_nodes(nodes) + # merged_nodes = merge_nodes(nodes, required_len) + merged_nodes = merge_nodes_recursion(nodes, required_len) + + cidr4s = [] + sum_added_ips = 0 + for ip_value, mask_len, added_ips, _ in merged_nodes: + cidr4s.append(make_cidr4(ip_value, mask_len)) + sum_added_ips += added_ips cidr4s_str = "\n".join(cidr4s) print( diff --git a/tests/test_cidr4_merger.py b/tests/test_cidr4_merger.py index 4035156..0be9164 100644 --- a/tests/test_cidr4_merger.py +++ b/tests/test_cidr4_merger.py @@ -2,15 +2,19 @@ import pytest from cidr4_merger import ( Cidr4MergerError, - answer, cidr4_to_node, + find_neighbours_singles, get_group_with_max_mask_len, get_net_addr, get_parent_ip, have_same_parent, + lift_lonely_node, make_cidr4, + make_groups, make_parent, + merge_neighbors, merge_nodes, + merge_nodes_recursion, reduce_nodes, sort_nodes, ) @@ -203,19 +207,3 @@ def test_merge_nodes(): ) assert exc_info.type is Cidr4MergerError assert str(exc_info.value) == "The top of the tree has no parent!" - - -def test_answer(): - assert answer( - [ - (0, 2, 12, 0), - (2147483648, 2, 1, 2147483648), - (3221225472, 2, 2, 2147483648), - ], - 2, - ) == (["128.0.0.0/1", "0.0.0.0/2"], 15) - - with pytest.raises(Exception) as exc_info: - answer([(0, 2, 0, 0), (2147483648, 2, 0, 2147483648)], 1) - assert exc_info.type is Cidr4MergerError - assert str(exc_info.value) == "The top of the tree has no parent!"