leetcode525. 连续数组 python

给定一个二进制数组, 找到含有相同数量的 0 和 1 的最长连续子数组。

示例 1:

输入: [0,1]
输出: 2
说明: [0, 1] 是具有相同数量0和1的最长连续子数组。

示例 2:

输入: [0,1,0]
输出: 2
说明: [0, 1] (或 [1, 0]) 是具有相同数量0和1的最长连续子数组

Python(超时)

class Solution(object):
    def findMaxLength(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        
        dict_nums = {0:0,1:0}
        max_len = 0
        len_nums = len(nums)
        for j in range(len_nums):
            for i in nums[j:]:
                if i == 0:
                    dict_nums[0] += 1
                elif i == 1:
                    dict_nums[1] += 1
                if dict_nums[0] == dict_nums[1]:
                    max_len = max(max_len,dict_nums[0]+dict_nums[1])
            dict_nums = {0: 0, 1: 0}
        return (max_len)

时间复杂度改为O(n),通过

class Solution(object):
    def findMaxLength(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        sums = [0] * len(nums)
        dmap = collections.defaultdict(int)
        last = 0
        for i, n in enumerate(nums):
            last += 2 * nums[i] - 1
            sums[i] = last
            dmap[last] = max(dmap[last], i)
        ans = 0
        for i, m in enumerate(sums):
            if m == 0:
                ans = max(ans, i + 1)
            else:
                ans = max(ans, dmap[m] - i)
        return ans