diff --git a/linked_lists/rotate_list_by_k.py b/linked_lists/rotate_list_by_k.py index 417c9c5..14ede5a 100644 --- a/linked_lists/rotate_list_by_k.py +++ b/linked_lists/rotate_list_by_k.py @@ -5,21 +5,23 @@ def rotate_list_by_k(head, k): - if head is None and head.next is None: - return None - - end, new_end, n = head, head, 1 - + if head is None: + return head + + # get the size of the list + end, n = head, 1 while end.next: end = end.next n += 1 - end.next = head - for i in range(n - k % n - 1): + # rotate + end.next = head + new_end, i = head, 0 + while i < n - (k % n) - 1: new_end = new_end.next - + i += 1 + + # remove cycle new_head = new_end.next - new_end.next = None return new_head -