Valid Subarrays With Matching Sum Digits I — Python

# Time:  O(nlogr)
# Space: O(n)

# prefix sum, two pointers
class Solution(object):
    def countValidSubarrays(self, nums, x):
        """
        :type nums: List[int]
        :type x: int
        :rtype: int
        """
        prefix = [0]*(len(nums)+1)
        for i in xrange(len(nums)):
            prefix[i+1] = prefix[i]+nums[i]
        result = 0
        base = 1
        while x*base <= prefix[-1]:
            cnt = [0]*10
            left = right = 0
            for i in xrange(len(nums)):
                while prefix[right] <= prefix[i+1]-x*base:
                    cnt[prefix[right]%10] += 1
                    right += 1
                while prefix[left] <= prefix[i+1]-(x+1)*base:
                    cnt[prefix[left]%10] -= 1
                    left += 1
                result += cnt[(prefix[i+1]-x)%10]
            base *= 10
        return result


# Time:  O(n^2 * logr)
# Space: O(1)
# brute force
class Solution2(object):
    def countValidSubarrays(self, nums, x):
        """
        :type nums: List[int]
        :type x: int
        :rtype: int
        """
        def check(n):
            if n%10 != x:
                return False
            while n//10:
                n //= 10
            return n == x

        result = 0
        for i in xrange(len(nums)):
            total = 0
            for j in xrange(i, len(nums)):
                total += nums[j]
                if check(total):
                    result += 1
        return result