class Solution(object):
    def findMaxLength(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """

        max_len = 0
        count_dict = {}
        count = 0
        index = 0
        for num in nums:
            if num == 0:
                count -= 1
            else:
                count += 1

            if count == 0 and index + 1 > max_len:
                max_len = index + 1

            old_index = count_dict.get(count)

            if old_index is None:
                count_dict[count] = index
            else:
                new_len = index - old_index
                if max_len < new_len:
                    max_len = new_len

            index += 1

        return max_len


sol = Solution()
print(sol.findMaxLength([0, 1]))