LeetCode 324. Wiggle Sort II

class Solution:
    def wiggleSort(self, nums: List[int]) -> None:
        n = len(nums)
        # Quick Select
        def findK(l, r, k):
            if l >= r:
                return nums[k]
            i = l - 1
            j = r + 1
            p = nums[(i+j)//2]
            while i < j:
                i += 1
                while nums[i] < p:
                    i += 1
                j -= 1
                while nums[j] > p:
                    j -= 1
                if i < j:
                    nums[i], nums[j] = nums[j], nums[i]
            if k <= j:
                return findK(l, j, k)
            return findK(j+1, r, k)
        """"
        Mapping
        [0, mid_index) -> [1, 3, 5, ...]
        [mid_index, n) -> [0, 2, 4, ...]
        1+2*i
        i=0 => 1
        i=1 => 3
        i=2 => 5
        i=3 => 7
        i=5 => 9
        1.n is even, nums[(1+2*i)%(n+1)]
          example: n = 4
          i=0 => 1
          i=1 => 3
          i=2 => 0
          i=3 => 2
        2.n is odd, nums[(1+2*i)%(n)]
          example: n = 3
          i=0 => 1
          i=1 => 0
          i=2 => 2
        1+2: nums[(1+2*i)%(n|1)]
        """
        A = lambda i: (1+2*i)%(n|1)
                    
        """
        Three-way parition to sort the array
        0-j: >mid, j-i:=mid, i-k:<mid
        Use the mapping to put in the right position
        """
        mid = findK(0, n-1, n//2)
        i = j = 0
        k = len(nums) - 1
        while i <= k:
            if nums[A(i)] > mid:
                nums[A(i)], nums[A(j)] = nums[A(j)], nums[A(i)]
                j += 1
                i += 1
            elif nums[A(i)] < mid:
                nums[A(i)], nums[A(k)] = nums[A(k)], nums[A(i)]
                k -= 1
            else:
                i += 1