Python, harder prefix sum approach (89ms)

  • 0

    Motivated by each triangle-stick having length <= 1000, let's record psum[s] = the number of sticks with length < s. We'll also remove pesky 0-length sticks as well.

    Now for every stick u and v, we want to know the number of sticks w occurring after v (if the sticks were in sorted order) with w < u+v. We can split this sum into ans += (number of sticks < u+v) and ans -= index of v + 1. The latter half can be calculated as (N-1)(N)(N+1)/3, so we only need to focus on the first half.

    With multiplicities, we can calculate this sum in a relatively straightforward way. We can also short circuit our loop: if we have say, stick lengths 10 and 15, and we know 20 is the largest stick, then we know that we will be able to make triangles freely for the rest of the middle sticks.

    def triangleNumber(self, A):
        B = max(A)
        if B == 0: return 0
        count = [0] * (B+1)
        N = 0
        for x in A:
            if x:
                count[x] += 1
                N += 1
        keys = [x for x in xrange(B+1) if count[x]]
        psum = [0]
        for c in count:
            psum.append(psum[-1] + c)
        #count[x] = the number of sticks of length x
        #psum[x] = the number of sticks < length x
        ans = 0
        for i1, k1 in enumerate(keys):
            for i2 in xrange(i1, len(keys)):
                k2 = keys[i2]
                #ans += multiplicity * (number of sticks < k1 + k2) 
                if i1 != i2:
                    multiplicity = count[k1] * count[k2]
                    multiplicity = count[k1] * (count[k1] - 1) / 2
                if k1 + k2 <= B:
                    ans += multiplicity * psum[k1 + k2]
                    multiplicity += count[k1] * (psum[-1] - psum[k2+1])
                    ans += multiplicity * N
        return ans - N*(N*N - 1)/3

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.