Segregate 0s and 1s in an array

• Time : $O(n)$
• Space : $O(1)$

vector<int> Solution::solve(vector<int> &A) {
    int l = 0, r = A.size() - 1;
    while(l < r) {
        while(A[l] == 0) l++;
        while(A[r] == 1) r--;
        if(l < r) swap(A[l],A[r]);
    }
    return A;
}