leetcode

commit 4b640063f1be1c56ae6aed07d3b50e345e518323
parent 356972e47558e0ff01e3d91952dacc303df61744
Author: Jack Sangdahl <jack@pngu.org>
Date:   Mon,  8 Jun 2026 13:58:07 -0600

2161. partition array according to given pivot

Diffstat:
M
partition_arr_pivot.c
 | 
64
++++++++++++++++++++++++++++++++++++++++++++++++----------------
M
partition_arr_pivot.py
 | 
9
++-------

2 files changed, 50 insertions(+), 23 deletions(-)
diff --git a/partition_arr_pivot.c b/partition_arr_pivot.c
@@ -11,32 +11,64 @@
  * rearrangement
  */
 
+uint64_t pack(int32_t a, int32_t b) {
+  return ((uint64_t)(b & 0xFFFFFFFF) << 32) | (a & 0xFFFFFFFF);
+}
+
+void unpack(uint64_t s, int32_t *a, int32_t *b) {
+  *a = (int32_t)(s & 0xFFFFFFFF);
+  *b = (int32_t)((s >> 32) & 0xFFFFFFFF);
+}
+
 int *pivotArray(int *nums, int numsSize, int pivot, int *returnSize) {
   int *ans = (int *)malloc(numsSize * sizeof(int));
-  int idx = 0;
-  for (int i = 0; i < numsSize; i++)
-    if (nums[i] < pivot)
-      ans[idx++] = nums[i];
-  for (int i = 0; i < numsSize; i++)
-    if (nums[i] == pivot)
-      ans[idx++] = nums[i];
-  for (int i = 0; i < numsSize; i++)
-    if (nums[i] > pivot)
-      ans[idx++] = nums[i];
   *returnSize = numsSize;
+  for (int i = 0; i < numsSize; i++)
+    ans[i] = nums[i];
+  int j = 0, cnt = 0;
+  for (int i = 0; i < numsSize; i++) {
+    if (ans[i] < pivot) {
+      if (ans[j] > pivot) {
+        uint64_t packed = pack(ans[i], ans[j]);
+        ans[j] = (int)packed;
+      } else
+        ans[j] = ans[i];
+      j++;
+    } else if (ans[i] == pivot)
+      cnt++;
+  }
+  j = numsSize - 1;
+  for (int i = numsSize - 1; i >= 0; i--) {
+    if (ans[i] > pivot) {
+      if (ans[i] > INT_MAX) {
+        int32_t a, b;
+        unpack((uint64_t)ans[i], &a, &b);
+        ans[j] = a;
+        ans[i] = b;
+      } else
+        ans[j] = ans[i];
+      j--;
+    }
+  }
+  for (int i = j + 1; i <= j + cnt; i++)
+    ans[i] = pivot;
   return ans;
 }
 
 int main() {
-  int n1[] = {9, 12, 5, 10, 14, 3, 10}, rs1;
-  int n2[] = {-3, 4, 3, 2}, rs2;
+  int n1[] = {9, 12, 5, 10, 14, 3, 10}, r1[] = {9, 5, 3, 10, 14, 3, 10}, rs1;
+  int n2[] = {-3, 4, 3, 2}, r2[] = {-3, 4, 2, 3}, rs2;
   int *pa1 = pivotArray(n1, ARRAY_SIZE(n1), 10, &rs1);
   int *pa2 = pivotArray(n2, ARRAY_SIZE(n2), 2, &rs2);
-  for (int i = 0; i < rs1; i++)
-    printf("%d ", pa1[i]); // expect: 9,5,3,10,10,12,14
+  for (int i = 0; i < rs1; i++) {
+    printf("%d ", pa1[i]);
+    assert(pa1[i] == r1[i]);
+  }
   printf("\n");
-  for (int i = 0; i < rs2; i++)
-    printf("%d ", pa2[i]); // expect: -3,2,4,3
+  for (int i = 0; i < rs2; i++) {
+    printf("%d ", pa2[i]);
+    assert(pa2[i] == r2[i]);
+  }
   printf("\n");
   free(pa1);
   free(pa2);
diff --git a/partition_arr_pivot.py b/partition_arr_pivot.py
@@ -11,13 +11,8 @@ rearrangement
 """
 
 
-class Solution(object):
-    def pivotArray(self, nums, pivot):
-        """
-        :type nums: List[int]
-        :type pivot: int
-        :rtype: List[int]
-        """
+class Solution:
+    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
         ans = [0] * len(nums)
         left, right = 0, len(nums) - 1
         for i, j in zip(range(len(nums)), range(len(nums) - 1, -1, -1)):