927. Three Equal Parts

Given an array A of 0s and 1s, divide the array into 3 non-empty parts such that all of these parts represent the same binary value.

If it is possible, return any [i, j] with i+1 < j, such that:

• A[0], A[1], ..., A[i] is the first part;
• A[i+1], A[i+2], ..., A[j-1] is the second part, and
• A[j], A[j+1], ..., A[A.length - 1] is the third part.
• All three parts have equal binary value.

If it is not possible, return [-1, -1].

Note that the entire part is used when considering what binary value it represents.  For example, [1,1,0] represents 6 in decimal, not 3.  Also, leading zeros are allowed, so [0,1,1] and [1,1] represent the same value.

 

Example 1:

Input: [1,0,1,0,1]
Output: [0,3]

Example 2:

Input: [1,1,0,1,1]
Output: [-1,-1]

 

Note:

1. 3 <= A.length <= 30000
2. A[i] == 0 or A[i] == 1

思路：two pointer，表示第一个数的起始位置i（设第一个数位p），第3个数的起始位置j（设第二个数位q），第2个数就是中间部分

如果p<q，i++

如果p>q，j--

如果p==q，j--，（不能i++，因为i++导致p变大，可能错过答案，比如j前面有0的情况）

class Solution:
    def threeEqualParts(self, a):
        """
        :type A: List[int]
        :rtype: List[int]
        """
        def comp(s1,s2):
            s1=s1.lstrip('0')
            s2=s2.lstrip('0')
            if len(s1)>len(s2): return 1
            elif len(s1)<len(s2): return -1
            elif s1>s2: return 1
            elif s1<s2: return -1
            return 0
        
        k=0
        while k<len(a) and a[k]==0: k+=1
        a=a[k:]
        if not a: return [0,k-1]
        
        p,q=str(a[0]),str(a[-1])
        i,j=1,len(a)-2
        m=''.join(list(map(str,a[1:-1])))
        if i==j: return comp(p,m)==0 and comp(q,m)==0
        if comp(p,m)==0 and comp(q,m)==0: return [k+i-1,k+j+1]
        while i<j:
            if comp(p,q)<0:
                p=p+str(a[i])
                m=m[1:]
                i+=1
            else:
                q=str(a[j])+q
                m=m[:-1]
                j-=1
            if comp(p,m)==0 and comp(q,m)==0: return [k+i-1,k+j+1]
        return [-1,-1] 
    
s=Solution()
print(s.threeEqualParts([0,1,1,0,1]))
print(s.threeEqualParts([0,0,0,0,0]))
print(s.threeEqualParts([1,1,0,0,1]))
print(s.threeEqualParts([1,0,1,0,1]))
print(s.threeEqualParts([1,1,0,1,1]))