Codewars kata - "Repetitive Sequence" I am trying to solve the following Codewars kata.
We are given a list as
seq = [0, 1, 2, 2]
We will have to write a function that will, firstly, add elements in the list using the following logic.
if n = 3 and as seq[n]=2, the new list will be seq = [0, 1, 2, 2, 3, 3]
if n = 4 and as seq[n]=3, the new list will be seq = [0, 1, 2, 2, 3, 3, 4, 4, 4]
if n = 5 and as seq[n]=3, the new list will be seq = [0, 1, 2, 2, 3, 3, 4, 4, 4, 5, 5, 5] and so on.
Then the function will return the n-th element of the list.
Some elements of the list:
[0, 1, 2, 2, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6, 6, 7, 7, 7, 7, 8, 8, 8, 8, 9, 9, 9, 9, 9,
10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13,
14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 16, 16, 16, 16, 16, 16, 16,
17, 17, 17, 17, 17, 17, 17, 18, 18, 18, 18, 18, 18, 18, 19, 19, 19, 19, 19, 19, 19,
20, 20, 20, 20, 20, 20, 20, 20, 21, 21, 21, 21, 21, 21, 21, 21]
Constraint for Python:
0 &lt;= n &lt;= 2^41
My code runs successfully in my system for any value of n, including n=2^41 (within 2.2s). But it times out in Codewars. Can anyone help me in optimizing my code? Thanks in advance.
My code:
def find(n):
    arr = [0, 1, 2, 2]
    if n &lt;= 3:
        return arr[n]
    else:
        arr_sum = 5
        for i in range(3, n+1):
            arr_sum += i * arr[i]
            if arr_sum &gt;= n:
                x = (arr_sum - n) // i
                return len(arr) + arr[i] - (x+1)
            else:
                arr += [i] * arr[i]
Building the list takes a lot of time. Note that your  arr += [i] * arr[i] repeats the same value i over and over again. This can be much compressed by just storing repeat(i, arr[i]) instead. This succeeds in about 6 seconds, well under their 12 seconds limit:
from itertools import chain, repeat

def find(n):
    if n &lt;= 3:
        return [0, 1, 2, 2][n]
    arr = [[2]]
    arr_sum = 5
    arr_len = 4
    for i, arr_i in enumerate(chain.from_iterable(arr), 3):
        arr_sum += i * arr_i
        if arr_sum &gt;= n:
            x = (arr_sum - n) // i
            return arr_len + arr_i - (x+1)
        arr.append(repeat(i, arr_i))
        arr_len += arr_i

Note that in the n &gt; 3 case, we start already with the number 3, appending it twice to the list. Thus of the sequence start [0, 1, 2, 2] we only need [2], so I start with arr = [[2]] (which is shorter than [repeat(2, 1)], and chain doesn't mind).

Alternatively... note that you're extending arr much faster than you're consuming it. For n=241, you grow it to over 51 million elements, but you're actually reading fewer than the first 70 thousand. So you could stop truly extending the list at that point. This succeeds in about 4.7 seconds:
def find(n):
    arr = [0, 1, 2, 2]
    if n &lt;= 3:
        return arr[n]
    arr_sum = 5
    arr_len = 4
    for i in range(3, n+1):
        arr_sum += i * arr[i]
        if arr_sum &gt;= n:
            x = (arr_sum - n) // i
            return arr_len + arr[i] - (x+1)
        arr_len += arr[i]
        if arr_len &lt; 70_000:
            arr += [i] * arr[i]


And... you can combine the above two improvements, i.e., apply that if arr_len &lt; 70_000: to the repeat-version. That then succeeds in about 4.5 seconds.
Benchmark results on my PC for n=241:
Your original:   1.795 seconds
My first one:    0.043 seconds (42 times faster)
My second one:   0.041 seconds (44 times faster)
The combination: 0.026 seconds (69 times faster)


Oh and a style comment: You twice do this:
if ...:
    return ...
else:
    ...

The else and the indentation of all the remaining code are unnecessary and I'd avoid it. I've done so in the above solutions.