Cloned undirected graph

leetcode/stdlib.py

@@ -147,3 +147,23 @@ class Node:
         ret += self.left.__list__() if self.left is not None else [None]
         ret += self.next.__list__() if self.next is not None else [None]
         return ret
+
+
+class NeighborlyNode:
+    """NeighborlyNode is a "Node" type used in leetcode graph puzzles"""
+
+    def __init__(self, val=0, neighbors=None):
+        self.val = val
+        self.neighbors = neighbors if neighbors is not None else []
+
+
+class NeighborlyNodeNicely(typing.NamedTuple):
+    val: int
+    neighbors: list["NeighborlyNodeNicely"]
+
+    def __eq__(self, other: typing.Optional["NeighborlyNodeNicely"]) -> bool:
+        return (
+            other is not None
+            and self.val == other.val
+            and [n.val for n in self.neighbors] == [n.val for n in other.neighbors]
+        )

leetcode/stuff.py

@@ -1,6 +1,5 @@
+import collections.abc
 import copy
-import itertools
-import math
 import random
 import typing
 
@@ -581,3 +580,48 @@ def accum_sub_array_maxes(nums: list[int]) -> list[int]:
         accum.append(max(nums[i], prev))
 
     return accum
+
+
+def neighborly_node_from_list(inlist: list[list[int]]):
+    # Alias "Node" type for leetcode compat
+    Node = stdlib.NeighborlyNodeNicely
+
+    if len(inlist) == 0:
+        return None
+
+    outlist = [Node(i + 1, []) for i in range(len(inlist))]
+
+    for i in range(len(inlist)):
+        outlist[i].neighbors[:] = []
+
+        for neighbor_val in inlist[i]:
+            outlist[i].neighbors.append(outlist[neighbor_val - 1])
+
+    return outlist[0]
+
+
+def neighborly_node_to_list(node) -> list[list[int]]:
+    serialized: dict[int, list[int]] = {}
+
+    for cur in traverse_neighborly_node(node, serialized):
+        if cur is None:
+            break
+
+        serialized[cur.val] = [n.val for n in cur.neighbors]
+
+    return [v for _, v in sorted(serialized.items())]
+
+
+def traverse_neighborly_node(
+    node: stdlib.NeighborlyNodeNicely, memo: collections.abc.Container[int]
+) -> typing.Iterator[stdlib.NeighborlyNodeNicely | None]:
+    yield node
+
+    if node is None:
+        return
+
+    for neighbor in node.neighbors:
+        if neighbor.val in memo:
+            continue
+
+        yield from traverse_neighborly_node(neighbor, memo)

leetcode/test_stuff.py

@@ -475,3 +475,19 @@ def test_trie_busy():
 )
 def test_max_sub_array(nums: list[int], expected: int):
     assert stuff.sum_max_sub_array(nums) == expected
+
+
+@pytest.mark.parametrize(
+    ("inlist",),
+    [
+        ([[2, 3, 4], [1, 7], [1], [1, 5, 6, 8], [4], [4], [2], [4]],),
+        ([[2, 4], [1, 3], [2, 4], [1, 3]],),
+    ],
+)
+def test_copy_neighborly_node(inlist):
+    orig = stuff.neighborly_node_from_list(inlist)
+    copied = stuff.neighborly_node_from_list(stuff.neighborly_node_to_list(orig))
+
+    assert id(orig) != id(copied)
+    assert orig == copied
+    assert stuff.neighborly_node_to_list(orig) == stuff.neighborly_node_to_list(copied)