Codesignal: extendedFibonacci

Original Problem

You want to impress your friends with the Fibonacci Sequence, but they all know it already! Instead, you decide to make Fibonacci-like sequences as follows:

Given nonnegative integers \(a\) and \(b\), the sequence \(\{s_n\}_{n\geq 0}\) is defined as:

\(s_0=2\)

\(s_1=a\)

\(s_{n+2}=as_{n+1}+bs_n\)

Since you want to impress your friends, you should be able to quickly tell them \(s_i\) for any reasonable \(i\).

This value might be so large that you can't easily display it. Instead, output the value mod \(10^9 + 7\) (which is prime), in the range \([0,10^9+7)\).

Note that it is guaranteed that \(a^2+4b\) is a perfect square, and that calculation of this value will not overflow a double.

Example

Input/Output

Solution

We want to finda closed form solution for the extended fibonacci sequence \(s_n=as_{n-1}+bs_{n-2}\). To do that we try to express the sequence as \(s_n=kx^n\), which means the sequence is \(kx^n=akx^{n-1}+bkx^{n-2}\). Simplifying this equation yields \(x^2-ax-b=0\), which has the solutions \(x=\frac{a\pm\sqrt{a^2+4b}}{2}\). The sequence now reads as

\[s_n=k_1\left(\frac{a+\sqrt{a^2+4b}}{2}\right)^n + k_2\left(\frac{a-\sqrt{a^2+4b}}{2}\right)^n\]

We can now use the base cases to solve for \(k_1, k_2\):

\[s_0=k_1\left(\frac{a+\sqrt{a^2+4b}}{2}\right)^0 + k_2\left(\frac{a-\sqrt{a^2+4b}}{2}\right)^0=2\]

\[s_1=k_1\left(\frac{a+\sqrt{a^2+4b}}{2}\right)^1 + k_2\left(\frac{a-\sqrt{a^2+4b}}{2}\right)^1 = a\]

From the first equation follows that \(k_1+k_2=2\) and using it to solve the second equation reveals that \(k_1=k_2=1\). It follows that the closed form solution under the modular congruence is:

\[s_n\equiv \left(\frac{a+\sqrt{a^2+4b}}{2}\right)^n + \left(\frac{a-\sqrt{a^2+4b}}{2}\right)^n\pmod{10^9+7}\]

Implementig the congruence is then quite simple:

M = 10**9 + 7
extendedFibonacci = lambda a, b, i: \
    sum(pow(int(a + k * (a * a + 4 * b)**.5) / 2, i, M) for k in (-1, 1)) % M

« Back to problem overview