Problem statement

https://leetcode.com/problems/n-th-tribonacci-number/

Solution 1

If you solved problem about Fibonacci numbers, this one will be a piece of cake for you. The idea is to use dynamic programming, where dp[i] is answer for i-th number. Then we can find dp[i] = dp[i-1] + dp[i-2] + dp[i-3]. Moreover, we do not need to use additional memory, because in fact we need to keep only last 3 numbers: a, b, c. Then we need to do n updates a, b, c = b, c, a + b + c. Note also, that I start with numbers a, b, c = dp[-2], dp[-1], dp[0], where dp[-1] = 0, because dp[-1] + dp[0] + dp[1] = dp[2] and dp[-2] = 1, because dp[-2] + dp[-1] + dp[0] = dp[1].

Complexity

Time complexity is O(n), space is O(1).

Code

class Solution:
    def tribonacci(self, n):
        a, b, c = 1, 0, 0
        for _ in range(n):
            a, b, c = b, c, a + b + c
            
        return c

Solution 2

If you plunge deep into solution of Fibonacci numbers problem, then you are aware that there is faster than O(n) solution for big n. The idea is so called fact matrix power. For every reccurunt equation we have, we can write it as evaluation of specific element of matrix of recurrence. In this problem M = [[1,1,1],[1,0,0],[0,1,0]], which is equivalent to a, b, c = b, c, a + b + c. Notice, that we can also do matrix power with given module MOD, which I keep in my code, but we do not need in in fact - I keep it here for more reusable code. Notice that we need to be careful with np implementation, do not use it if M = 10**9 + 7 and M > 8, see my remark here https://flykiller.github.io/patterns/linear%20algebra

Complexity

Time complexity is O(log n * M^3), where M = 3 is the size of matrix. For given constraits there will be no gain in speed, but for bigger n there will be. Space complexity is O(M^3).

Code

import numpy as np

class Solution:
    def tribonacci(self, n):
        def power(M, n, MOD):
            result = np.eye(len(M), dtype = np.int64)
            while n > 0:
                if n%2: result = np.dot(M, result) % MOD
                M = np.dot(M, M) % MOD
                n //= 2
            return result
        
        t = power(np.array([[1,1,1],[1,0,0],[0,1,0]]), n, 1<<32)
        return t[0][-1]