diff --git a/leetcode/stdlib.py b/leetcode/stdlib.py index 20f67f1..3c19d17 100644 --- a/leetcode/stdlib.py +++ b/leetcode/stdlib.py @@ -1,3 +1,4 @@ +import dataclasses import typing @@ -14,6 +15,30 @@ class ListNode: self.next = next +class ListNodeRandom: + """ListNodeRandom is another weirdo linked list thing from + leetcode that is obviously very different than a binary tree + node :upside_down_face: + """ + + def __init__( + self, + x: int, + next: typing.Optional["ListNodeRandom"] = None, + random: typing.Optional["ListNodeRandom"] = None, + ): + self.val = x + self.next = next + self.random = random + + +@dataclasses.dataclass +class ListNodeRandomNicely: + val: int + next: typing.Optional["ListNodeRandomNicely"] = None + random: typing.Optional["ListNodeRandomNicely"] = None + + class BinaryTreeNode(typing.Protocol): val: int left: typing.Optional["BinaryTreeNode"] diff --git a/leetcode/stuff.py b/leetcode/stuff.py index 835b6b9..3fc6a18 100644 --- a/leetcode/stuff.py +++ b/leetcode/stuff.py @@ -509,3 +509,29 @@ def count_factorial_trailing_zeroes(number: int) -> int: divisor *= 5 return zeroes_count + + +def copy_random_list( + head: stdlib.ListNodeRandom | None, +) -> stdlib.ListNodeRandom | None: + if head is None: + return None + + ordered = [] + cur = head + while cur is not None: + ordered.append(cur) + cur = cur.next + + ordered_copy = [stdlib.ListNodeRandom(entry.val) for entry in ordered] + + hash_idx = {hash(n): i for i, n in enumerate(ordered)} + + for i, entry in enumerate(ordered): + if i + 1 < len(ordered_copy): + ordered_copy[i].next = ordered_copy[i + 1] + + if entry.random is not None: + ordered_copy[i].random = ordered_copy[hash_idx[hash(entry.random)]] + + return ordered_copy[0]