add reverse k group linkedlist

data_structures/linked_list/reverse_k_group.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+@dataclass
+class ListNode:
+    data: int
+    next_node: ListNode | None = None
+
+def print_linked_list(head: ListNode | None) -> None:
+    """
+    Print the entire linked list iteratively.
+
+    This function prints the elements of a linked list separated by '->'.
+
+    Parameters:
+        head (ListNode | None): The head of the linked list to be printed,
+        or None if the linked list is empty.
+
+    >>> head = insert_node(None, 0)
+    >>> head = insert_node(head, 2)
+    >>> head = insert_node(head, 1)
+    >>> print_linked_list(head)
+    0->2->1
+    >>> head = insert_node(head, 4)
+    >>> head = insert_node(head, 5)
+    >>> print_linked_list(head)
+    0->2->1->4->5
+    """
+    if head is None:
+        return
+    while head.next_node is not None:
+        print(head.data, end="->")
+        head = head.next_node
+    print(head.data)
+
+def insert_node(head: ListNode | None, data: int) -> ListNode:
+    """
+    Insert a new node at the end of a linked list and return the new head.
+
+    Parameters:
+        head: The head of the linked list.
+        data: The data to be inserted into the new node.
+
+    Returns:
+        The new head of the linked list.
+
+    >>> head = insert_node(None, 10)
+    >>> head = insert_node(head, 9)
+    >>> head = insert_node(head, 8)
+    >>> print_linked_list(head)
+    10->9->8
+    """
+    new_node = ListNode(data)
+    if head is None:
+        return new_node
+
+    temp_node = head
+    while temp_node.next_node:
+        temp_node = temp_node.next_node
+
+    temp_node.next_node = new_node
+    return head
+
+class Solution:
+    def reverse_k_group(self, head: ListNode | None, k: int) -> ListNode | None:
+        """
+        Reverse k-sized groups of nodes in a linked list.
+
+        Parameters:
+            head: The head of the linked list.
+            k: The size of each group to reverse.
+
+        Returns:
+            The head of the reversed linked list.
+
+        >>> head = insert_node(None, 1)
+        >>> head = insert_node(head, 2)
+        >>> head = insert_node(head, 3)
+        >>> head = insert_node(head, 4)
+        >>> head = insert_node(head, 5)
+        >>> head = insert_node(head, 6)
+        >>> solution = Solution()
+        >>> new_head = solution.reverse_k_group(head, 2)
+        >>> print_linked_list(new_head)
+        2->1->4->3->6->5
+        """
+        def reverse_group(head: ListNode | None, k: int) -> tuple:
+            prev_group_tail = None
+            nodes_left = k
+            current_group_head = head
+
+            while current_group_head:
+                current_group_head = current_group_head.next_node
+                nodes_left -= 1
+
+            if nodes_left > 0:
+                return head, None, None, False
+
+            current_tail = head
+
+            while head and k > 0:
+                k -= 1
+                next_node = head.next_node
+                head.next_node = prev_group_tail
+                prev_group_tail = head
+                head = next_node
+
+            return prev_group_tail, current_tail, head, True
+
+        new_head, current_tail, next_group_head, success = reverse_group(head, k)
+
+        while success:
+            new_group_head, new_group_tail, next_next_group_head, success = reverse_group(next_group_head, k)
+            current_tail.next_node = new_group_head
+            current_tail = new_group_tail
+            next_group_head = next_next_group_head
+
+        return new_head
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
+    head = insert_node(None, 5)
+    head = insert_node(head, 1)
+    head = insert_node(head, 2)
+    head = insert_node(head, 4)
+    head = insert_node(head, 3)
+
+    print("Original list: ", end="")
+    print_linked_list(head)
+
+    k = 3
+    solution = Solution()
+    new_head = solution.reverse_k_group(head, k)
+
+    print(f"After reversing groups of size {k}: ", end="")
+    print_linked_list(new_head)