diff --git a/2161. Partition Array According to Given Pivot b/2161. Partition Array According to Given Pivot new file mode 100644 index 0000000..5231d66 --- /dev/null +++ b/2161. Partition Array According to Given Pivot @@ -0,0 +1,36 @@ +class Solution { +public: + vector pivotArray(vector& nums, int pivot) { + int n = nums.size(); + vector ans(n); + int lesser = 0, greater = 0, count = 0; + int lesserIndex = 0, pivotIndex = 0, greaterIndex = 0; + + for (int i = 0; i < n; i++) { + if (nums[i] < pivot) { + lesser++; + } + else if (nums[i] == pivot) { + count++; + } + else { + greater++; + } + } + pivotIndex = lesser; + greaterIndex = lesser + count; + + for (int index = 0; index < n; index++) { + if (lesserIndex < lesser && nums[index] < pivot) { + ans[lesserIndex++] = nums[index]; + } + else if (pivotIndex < (lesser + count) && nums[index] == pivot) { + ans[pivotIndex++] = nums[index]; + } + else if (greaterIndex < n && nums[index] > pivot) { + ans[greaterIndex++] = nums[index]; + } + } + return ans; + } +};