Compress challenge

compress.py

"""
Given an array of characters with repeats, compress it in place.
The length after compression should be less than or equal to the original array.

Example:
Input: ['a', 'a', 'b', 'c', 'c', 'c']
Output: ['a', '2', 'b', 'c', '3']
"""
import unittest


def get_count_dict(char_list):
    char_count = {}

    for char in char_list:
        if char in char_count:
            char_count[char] += 1
        else:
            char_count[char] = 1

    return char_count


def get_count_list(char_count):
    char_count_list = []

    for key, val in char_count.items():
        if val > 1:
            char_count_list += [key, str(val)]
        else:
            char_count_list += [key]

    return char_count_list


def compress(char_list):
    return get_count_list(
        get_count_dict(char_list)
    )


class TestCompress(unittest.TestCase):

    def test_compress(self):
        compressed_array = compress(['a', 'a', 'b', 'c', 'c', 'c'])
        self.assertEqual(compressed_array, ['a', '2', 'b', 'c', '3'])


if __name__ == "__main__":
    unittest.main()