partition_element = 4 # last element

[3, 5, 1, 2, 6, 4]

^^

[3, 5, 1, 2, 6, 4] # swap

^^

[3, 5, 1, 2, 6, 4]

^  ^

[3, 1, 5, 2, 6, 4]

   ^  ^

[3, 1, 2, 5, 6, 4]

      ^  ^

swap element at len(nums) - 1 with element at i

[3, 1, 2, 4, 6, 5]

# NOTE: end is inclusive
def partition(nums, start, end):
  # Use element at end as partition element
  i = start
  for j in range(start, end):
    if nums[j] <= nums[end]:
      temp = nums[i]
      nums[i] = nums[j]
      nums[j] = temp
      i += 1
  temp = nums[i]
  nums[i] = nums[end]
  nums[end] = nums[i]
  # nums[i] is in same place as sorted array
  return nums, i

Iterate for start to end

Keep an extra index, i, starting at start

If element at index j is less than partition element:

Swap element at index j with element at index i

Increment i -> when iteration is over, all element to left of i is smaller than the partition element

Note we go from range(start, end) because element at end is the partition element, so we don’t need to check it

Swap element at index end with element at index i