Leetcode 443. String Compression

class Solution {
public:
    int compress(vector<char>& chars) {
        if(chars.size() == 0 || chars.size() == 1) {
            return chars.size();
        }
        int count = 1;
        int index = 0;
        char current = chars[0];
        for(int i = 1; i < chars.size(); i++) {
            if(chars[i] != current) {
                chars[index] = current;
                index++;
                if(count > 1) {
                    string temp = to_string(count);
                    for(char ch : temp) {
                        chars[index] = ch;
                        index++;
                    }
                }
                count = 1;
                current = chars[i];
            }
            else {
                count++;
            }
        }
        chars[index] = current;
        index++;
        if(count > 1) {
            for(char ch : to_string(count)) {
                chars[index] = ch;
                index++;
            }
        }
        return index;
    }
};