The partition step is similar to what we had in the quicksort algorithm. There are a couple of things that are worth noting:

    def partition(unsorted_array, first_index, last_index): 
        if first_index == last_index: 
            return first_index 

        pivot = unsorted_array[first_index] 
        pivot_index = first_index 
        index_of_last_element = last_index 

        less_than_pivot_index = index_of_last_element 
        greater_than_pivot_index = first_index + 1 

        while True: 

            while unsorted_array[greater_than_pivot_index] < pivot and  
                  greater_than_pivot_index < last_index: 
                  greater_than_pivot_index += 1 
            while unsorted_array[less_than_pivot_index] > pivot and 
                  less_than_pivot_index >= first_index: 
                  less_than_pivot_index -= 1 

            if greater_than_pivot_index < less_than_pivot_index: 
                temp = unsorted_array[greater_than_pivot_index] 
                unsorted_array[greater_than_pivot_index] = 
                    unsorted_array[less_than_pivot_index] 
                unsorted_array[less_than_pivot_index] = temp 
            else: 
                break 

        unsorted_array[pivot_index] =  
            unsorted_array[less_than_pivot_index] 
        unsorted_array[less_than_pivot_index] = pivot 

        return less_than_pivot_index 

An if statement has been inserted at the beginning of the function definition to cater for situations where first_index is equal to last_index. In such cases, it means that there is only one element in our sublist. Therefore, we simply return any of the function parameters, in this case, first_index.

The first element is always chosen as the pivot. This choice to make the first element the pivot is a random decision. It often does not yield a good split—and subsequently—a good partition. However, the i^th element will eventually be found, even though the pivot is chosen at random.

The partition function returns the pivot index pointed to by less_than_pivot_index, as we saw in the preceding chapter.