sort01 added using partition method of quick sort

Author: komal.bansal1

Array/sort_0_1.py

@@ -0,0 +1,35 @@
+'''
+https://www.geeksforgeeks.org/problems/segregate-0s-and-1s5106/1
+Given an array of length n consisting of only 0's and 1's in random order. Modify the array in-place to segregate 0s on the left side and 1s on the right side of the array.
+
+Example 1:
+Input:
+n = 5
+arr[] = {0, 0, 1, 1, 0}
+Output: {0, 0, 0, 1, 1}
+Explanation: 
+After segregate all 0's on the left and 1's on the right modify array will be {0, 0, 0, 1, 1}.
+
+Example 2:
+Input:
+n = 4
+arr[] = {1, 1, 1, 1}
+Output: {1, 1, 1, 1}
+Explanation: There are no 0 in the given array, so the modified array is 1 1 1 1.
+'''
+
+# use partition method of quick sort to move 0 on left and 1 on right side
+class Solution:
+    def segregate0and1(self, arr, n):
+        
+        pivot = 0
+        i = 0
+        j = 0
+        
+        while j < n:
+            if arr[j] == pivot:
+                arr[i], arr[j] = arr[j], arr[i]
+                i += 1
+                j += 1
+            else:
+                j += 1

Array/sort_0_1_2.py

@@ -22,22 +22,25 @@
 # if mid = 0, swap with low then low += 1 & mid += 1
 # if mid = 1, mid += 1
 # else swap mid with high, high -= 1
+# similar to sort_0_1, following partitioning method
+# i = 0 region, j = 1 region, k = 2 region
 class Solution:
-    def sortColors(self, arr: List[int]):
-    
-        n = len(arr)
-        low = 0
-        high = n-1
-        mid = 0
-        
-        while mid <= high:
-            
-            if arr[mid] == 0:
-                arr[mid], arr[low] = arr[low], arr[mid]
-                low += 1
-                mid += 1
-            elif arr[mid] == 1:
-                mid += 1
-            else:
-                arr[mid], arr[high] = arr[high], arr[mid]
-                high -= 1
+    def sortColors(self, nums: List[int]) -> None:
+        """
+        Do not return anything, modify nums in-place instead.
+        """
+        i = 0
+        j = 0
+        k = len(nums)-1
+
+        while j <= k:
+            if nums[j] == 1:    # if j is in correct region
+                j += 1
+            elif nums[j] == 0:  # if j found 0, means move to i region
+                nums[i], nums[j] = nums[j], nums[i]
+                i += 1
+                j += 1
+            else:               # if j found 2, means move to k region
+                nums[k], nums[j] = nums[j], nums[k]
+                k -= 1
+