Update rotate_list_by_k.py

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
-