class Solution(object): def checkRecord(self, n): ans=(1,0,0,0,0,0) md = 10**9+7 m=((1,0,1,0,1,0),(1,1,1,1,1,1),(1,0,0,0,0,0),(0,1,0,0,0,0),(0,0,1,0,0,0),(0,0,0,1,0,0)) def matrixmult(a,b): return tuple(tuple(sum(row[i]*b[i][column] for i in range(len(a)))%md for column in range(len(b))) for row in a) def matrixmult2(a,bv): return tuple(sum(row[i]*bv[i] for i in range(len(a)))%md for row in a) if n%2: ans = matrixmult2(m,ans) n/=2 while n>0: m=matrixmult(m,m) if n%2: ans = matrixmult2(m,ans) n/=2 return sum(ans)%md
The ans vector denotes different types of answer:
0 : no A, no ending L
1 : has A, no ending L
2 : no A, 1 ending L (i.e. ending in "PL")
3 : has A, 1 ending L (i.e. ending in "PL" or "AL")
4 : no A, 2 ending L (i.e. ending in "PLL")
5 : has A, 2 ending L (i.e. ending in "PLL" or "ALL")
The matrix m initially denotes appending a digit to an existing string:
type 0: P->0, A->1, L->2
type 1: P->1, L->3
type 2: P->0, A->1, L->4
type 3: P->1, L->5
type 4: P->0, A->1
type 5: P->1
As we only have 1 string, "" (of type 0), for N = 0. The answer matrix is initialized as (1,0,0,0,0,0).
With exponentiation by squaring, suitable matrices are multiplied to the answer vector as soon as they're found.
This works excellent for very large Ns, but for the test cases, an O(N) method is easily devised and turn out to perform the most quickly. This method is still not bad though.
The matrix multiplication functions can also be reused in other cases, yet it doesn't check whether the sizes match, nor that the inputs are valid.
P.S. I made a mistake SO dumb, to store all helper matrices in an array in the original post. It beats 85.71%; still pretty fast.