Codesignal Solution: 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

• For a = 1, b = 2 and i = 4, the output should be
  extendedFibonacci(a, b, i) = 17.
  s₀ = 2
  s₁ = a = 1
  s_{n + 2} = as_{n + 1} + bs_n = s_{n + 1} + 2s_n.
  s₂ = 1 + 2 · 2 = 5
  s₃ = 5 + 2 · 1 = 7
  s₄ = 7 + 2 · 5 = 17

• For a = 0, b = 15129 and i = 2831, the output should be
  extendedFibonacci(a, b, i) = 0.
  s₁ = a = 0
  s_{n + 2} = as_{n + 1} + bs_n = 15129s_n
  s₃ = 15129 · 0 = 0
  s₅ = 15129 · 0 = 0
  ...
  s₂₈₃₁ = 15129 · 0 = 0

Input/Output

• [time limit] 6000ms (groovy)

• [input] integer64 a

  Along with \(b\), defines the sequence \(\{s_n\}_{n\geq 0}\)
  \(a^2+4b\) is a perfect square.
  \(0\leq a\leq 2^{26}\)

• [input] integer64 b

  Along with \(a\), defines the sequence \(\{s_n\}_{n\geq 0}\)
  \(a^2+4b\) is a perfect square.
  \(0\leq b\leq 2^{26}\)

• [input] integer64 i

  Index of the desired value in the sequence
  \(0\leq i\leq 2^{53}\)

• [output] integer

  \(s_i \mod 10^9+7\)

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