Number of ways to remove items from box

Question

I encountered the following algorithmic question which has strict constraints on runtime (<10s and no large memory footprint) and I am stumped. My approach fails half the test cases.

Question

A box contains a number of items that can only br removed 1 at a time or 3 at a time.

How many ways can the box be emptied? the answer can be very large so return it as modulo of 10^9+7.

for example,there are n=7 items initially.They can be removed nine ways,as follows:

1.(1,1,1,1,1,1,1)
2.(1.1.1.1.3)
3.(1,1,1,3,1)
4.(1,1,3,1,1)
5.(1,3,1,1,1)
6.(3,1,1,1,1)
7.(1,3,3)
8.(3,1,3)
9.(3,3,1)

So the function should return 9.

Function Description: Your function must take in a parameter, n for the number of items, and return an integer which denotes the number of ways to empty the box.

Constraints: 1<=n<=10^8

Sample cases :

Input: 1
Sample OutPut: 1
Explanation: There is only 1 way to remove 1 item. Answer=(1%1000000007)=1

Input: 7
Sample OutPut: 9
There is only 9 ways to remove 7 items

My Approach

This leads to a standard recurrence relation where f(n) = f(n-3) + f(n-1) for n > 2, so i did it as follows

def memoized_number_of_ways(dic, n):
    if n not in dic:
        dic[n] = memoized_number_of_ways(dic, n-3) + memoized_number_of_ways(dic, n-1)
    return dic[n]

def numberOfWays(n):
    # Write your code here
    memoize = {1:1,2:1,3:2}
    import math
    ans = memoized_number_of_ways(memoize,n)
    return ans % (math.pow(10,9) + 7)

However this fails on any case where n > 10**2. How can you do this problem while accomodating n up to 10^8 and in less than 10s with not much memory?

Answer 1

Just write your recurrence using matrices (pardon my way of writing matrices, StackOverflow doesn't allow LaTeX).

[f(n)  ] = [1 0 1] [f(n-1)  ]
[f(n-1)] = [1 0 0] [f(n-2)  ]
[f(n-2)] = [0 1 0] [f(n-3)  ]

Now all you have to do is raise a 3x3 matrix (modulo fixed constant) to power n (or n-3 or something like that, depending on your "base case column vector", fill in the details), and then multiply it by a "base case column vector". This can be done in time O(logn).

PS: You may want to lookup matrix exponentiation.

Answer 2

Three solutions, fastest takes about 31 μs for n=10⁸ on tio.run (which has medium-fast computers).

A matrix power solution like described by advocateofnone that takes about 1 millisecond (Try it online!):

import numpy as np
from time import time

class ModInt:
    def __init__(self, x):
        self.x = x % (10**9 + 7)
    def __add__(a, b):
        return ModInt(a.x + b.x)
    def __mul__(a, b):
        return ModInt(a.x * b.x)
    def __str__(self):
        return str(self.x)

def solve(n):
    O = ModInt(0)
    I = ModInt(1)
    A = np.matrix([[O,I,O], [O,O,I], [I,O,I]])
    return (A**n)[2,2]

for _ in range(3):
    t0 = time()
    print(solve(10**8), time() - t0)

Output (result and time in seconds for n=10⁸, three attempts):

109786077 0.0010712146759033203
109786077 0.0010180473327636719
109786077 0.0009677410125732422

Another, taking about 0.5 milliseconds (Try it online!):

import numpy as np
from time import time

def solve(n):
    A = np.matrix([[0,1,0], [0,0,1], [1,0,1]])
    power = 1
    mod = 10**9 + 7
    while n:
        if n % 2:
            power = power * A % mod
        A = A**2 % mod
        n //= 2
    return power[2,2]

for _ in range(3):
    t0 = time()
    print(solve(10**8), time() - t0)

One based on @rici's solution in the comments, takes about 31 μs (Try it online!):

from timeit import repeat

def solve(n):
    m = 10**9 + 7
    def abc(n):
        if n == 0:
            return 0, 1, 0
        a, b, c = abc(n // 2)
        d = a + c
        e = b + d
        A = 2*a*b + c*c
        C = 2*b*c + d*d
        E = 2*c*d + e*e
        D = A + C
        B = E - D
        if n % 2:
            A, B, C = B, C, D
        return A%m, B%m, C%m
    return sum(abc(n)) % m

n = 10**8
print(solve(n))
for _ in range(3):
    t = min(repeat(lambda: solve(n), 'gc.enable()', number=1000)) / 1000
    print('%.1f μs' % (t * 1e6))

Explanation: Looking at the matrix powers from my previous solutions, I noticed they only actually contain five different values, and they're consecutive result numbers from our desired sequence. For example, A**19 is:

[[277 189 406]
 [406 277 595]
 [595 406 872]]

I gave them names in increasing order:

| b a c |
| c b d |
| d c e |

Squaring that matrix results in a matrix for larger n, with entries A/B/C/D/E. And if you square the above matrix, you'll find the relationships A = 2*a*b + c*c etc.

My helper function abc(n) computes the entries a/b/c of the n-th matrix power. For n=0, that's the identity matrix, so my a/b/c are 0/1/0 there. And in the end I, return the e-value (computed as e=b+d=a+b+c).