Write a program to solve the Dutch National Flag problem.
Implement def sort_colors(nums): to rearrange an array containing only 0, 1, and 2 so equal values are adjacent and appear in ascending order. Modify the input list in place and return it. Aim for linear time and constant extra space.
def sort_colors(nums):