add parent_ip info to Node tuple

This commit is contained in:
Pavel Patsey
2025-01-18 20:15:21 +03:00
parent edcce93d01
commit b5a1498a51
2 changed files with 146 additions and 111 deletions
+36 -31
View File
@@ -1,7 +1,10 @@
import cProfile
from typing import Optional
Node = tuple[int, int, int]
Node = tuple[int, int, int, int]
class Cidr4MergerError(Exception):
pass
def get_data(input_file):
@@ -11,12 +14,13 @@ def get_data(input_file):
def cidr4_to_node(cidr4: str) -> Node:
ip, mask_len = cidr4.strip().split("/")
ip_str, mask_len = cidr4.strip().split("/")
mask_len = int(mask_len)
a, b, c, d = list(map(int, ip.split(".")))
ip_value = a * 256**3 + b * 256**2 + c * 256**1 + d * 256**0
a, b, c, d = list(map(int, ip_str.split(".")))
ip = a * 256**3 + b * 256**2 + c * 256**1 + d * 256**0
added_ips_number = 0
return ip_value, mask_len, added_ips_number
parent_ip = get_parent_ip(ip, mask_len)
return ip, mask_len, added_ips_number, parent_ip
def sort_nodes(nodes: list[Node]) -> list[Node]:
@@ -29,16 +33,14 @@ def get_net_addr(ip: int, mask_len: int) -> int:
return net_addr
def get_parent_mask(ip: int, mask_len: int) -> Optional[int]:
def get_parent_ip(ip: int, mask_len: int) -> int:
if mask_len == 0:
return None
raise Cidr4MergerError("The top of the tree has no parent!")
return get_net_addr(ip, mask_len - 1)
def have_same_parent(ip_a, mask_len_a, ip_b, mask_len_b) -> bool:
return mask_len_a == mask_len_b and get_parent_mask(
ip_a, mask_len_a
) == get_parent_mask(ip_b, mask_len_b)
def have_same_parent(mask_len_a, parent_ip_a, mask_len_b, parent_ip_b) -> bool:
return mask_len_a == mask_len_b and parent_ip_a == parent_ip_b
def get_group_with_max_mask_len(nodes: list[Node]) -> list[Node]:
@@ -46,16 +48,19 @@ def get_group_with_max_mask_len(nodes: list[Node]) -> list[Node]:
return list(filter(lambda x: x[1] == max_mask_len, nodes))
def get_parent(a: Node, b: Optional[Node] = None) -> Node:
ip_a, mask_len_a, added_ips_number_a = a
ip = get_parent_mask(ip_a, mask_len_a)
mask_len = mask_len_a - 1
added_ips_number = added_ips_number_a + 2 ** (32 - mask_len_a)
def make_parent(a: Node, b: Node | None = None) -> Node:
ip_a, mask_len_a, added_ips_number_a, parent_ip_a = a
if b:
ip_b, mask_len_b, added_ips_b = b
assert have_same_parent(ip_a, mask_len_a, ip_b, mask_len_b)
ip_b, mask_len_b, added_ips_b, parent_ip_b = b
if not have_same_parent(mask_len_a, parent_ip_a, mask_len_b, parent_ip_b):
raise Cidr4MergerError("Nodes must be neighbors!")
added_ips_number = added_ips_number_a + added_ips_b
return ip, mask_len, added_ips_number
else:
added_ips_number = added_ips_number_a + 2 ** (32 - mask_len_a)
ip = parent_ip_a
mask_len = mask_len_a - 1
parent_ip = get_parent_ip(ip, mask_len)
return ip, mask_len, added_ips_number, parent_ip
def reduce_nodes(nodes: list[Node]) -> list[Node]:
@@ -66,9 +71,9 @@ def reduce_nodes(nodes: list[Node]) -> list[Node]:
i = 0
while i < len(group) - 1:
a, b = group[i], group[i + 1]
ip_a, mask_len_a, _ = a
ip_b, mask_len_b, _ = b
if have_same_parent(ip_a, mask_len_a, ip_b, mask_len_b):
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:
@@ -78,14 +83,14 @@ def reduce_nodes(nodes: list[Node]) -> list[Node]:
loners.append(group[i])
if neighbours:
zipped = zip(neighbours, map(lambda x: get_parent(x[0], x[1]), neighbours))
zipped = zip(neighbours, map(lambda x: make_parent(x[0], x[1]), neighbours))
min_zipped = min(zipped, key=lambda x: x[1][2])
(a, b), parent = min_zipped
nodes.remove(a)
nodes.remove(b)
nodes.append(parent)
elif loners:
zipped = zip(loners, map(get_parent, loners))
zipped = zip(loners, map(make_parent, loners))
min_zipped = min(zipped, key=lambda x: x[1][2])
a, parent = min_zipped
nodes.remove(a)
@@ -102,10 +107,10 @@ def merge_nodes(nodes: list[Node], required_len: int) -> list[Node]:
return nodes
def node_to_cidr4(ip_value, mask_len) -> str:
lst = [str(ip_value >> (i << 3) & 0xFF) for i in reversed(range(4))]
ip = ".".join(lst)
return f"{ip}/{mask_len}"
def make_cidr4(ip, mask_len) -> str:
lst = [str(ip >> (i << 3) & 0xFF) for i in reversed(range(4))]
ip_str = ".".join(lst)
return f"{ip_str}/{mask_len}"
def answer(nodes: list[Node], required_len: int) -> tuple[list[str], int]:
@@ -114,8 +119,8 @@ def answer(nodes: list[Node], required_len: int) -> tuple[list[str], int]:
cidr4s = []
sum_added_ips = 0
for ip_value, mask_len, added_ips in merged_nodes:
cidr4s.append(node_to_cidr4(ip_value, mask_len))
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