How to Sort linked list of 0s, 1s and 2s

[Question]: Given a singly linked list head sort it with with all with zeros, ones and two’s
Example:
Input: 1 -> 1 -> 2 -> 0 -> 2 -> 0 -> 1 -> NULL 
Output: 0 -> 0 -> 1 -> 1 -> 1 -> 2 -> 2 -> NULL

Approach: Create count array which stores count of 0’s, 1’s and 2’s occurrences. We will iterate over linked list & finds the count of each(0,1,2)
Create a new head named `current` and replace items with sorted order

// TC: O(n)
// SC: O(n)
func sortList012(_ head: ListNode?) {
    var count = [ 0, 0, 0 ]
    // Count the number of 0's, 1's, and 2's in the
    // linked list
    var current = head
    while current != nil {
        count[current!.val]+=1// Filling respective index like count of 0's 1's & 2's
        current = current?.next
    }
    // Overwrite the linked list with the sorted
    // elements
    current = head
    var i = 0
    while current != nil {
        if count[i] == 0 {
            i+=1
        } else {
            current?.val = i
            count[i]-=1
            current = current?.next
        }
    }
}

let zeroOneTwoLinkedList = ListNode(1, ListNode(2, ListNode(0,ListNode(2, ListNode(0)))))
sortList012(zeroOneTwoLinkedList)
//print("sorted by 0,1,2", zeroOneTwoLinkedList)// 0 0 1 2 1