Update Split Linked List in Parts.java

varunu28 · web-flow · commit 5642e0cadbee · 2023-09-06T07:35:18.000-07:00
diff --git a/Medium/Split Linked List in Parts.java b/Medium/Split Linked List in Parts.java
@@ -7,25 +7,27 @@
  *     ListNode(int val) { this.val = val; }
  *     ListNode(int val, ListNode next) { this.val = val; this.next = next; }
  * }
-*/
+ */
 class Solution {
-  public ListNode[] splitListToParts(ListNode head, int k) {
-    ListNode[] arr = new ListNode[k];
-    int nodeLength = 0;
-    for (ListNode curr = head; curr != null; curr = curr.next) {
-      nodeLength++;
+    public ListNode[] splitListToParts(ListNode head, int k) {
+        int length = 0;
+        ListNode curr = head;
+        while (curr != null) {
+            length++;
+            curr = curr.next;
+        }
+        ListNode[] splits = new ListNode[k];
+        int elementsPerPart = length / k;
+        int partsWithExtra = length % k;
+        ListNode prev = null;
+        for (int i = 0; i < k && head != null; i++, partsWithExtra--) {
+            splits[i] = head;
+            for (int j = 0; j < elementsPerPart + (partsWithExtra > 0 ? 1 : 0); j++) {
+                prev = head;
+                head = head.next;
+            }
+            prev.next = null;
+        }
+        return splits;
     }
-    int n = nodeLength / k;
-    int remaining = nodeLength % k;
-    ListNode prev = null;
-    for (int i = 0; i < k && head != null; i++, remaining--) {
-      arr[i] = head;
-      for (int j = 0; j < n + (remaining > 0 ? 1 : 0); j++) {
-        prev = head;
-        head = head.next;
-      }
-      prev.next = null;
-    }
-    return arr;
-  }
 }
