4
\$\begingroup\$

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 <= n <= 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 <= 3:
        return arr[n]
    else:
        arr_sum = 5
        for i in range(3, n+1):
            arr_sum += i * arr[i]
            if arr_sum >= n:
                x = (arr_sum - n) // i
                return len(arr) + arr[i] - (x+1)
            else:
                arr += [i] * arr[i]
\$\endgroup\$
5
  • \$\begingroup\$ I think your best bet is deriving a mathematical formula to calculate nth element of the sequence. \$\endgroup\$ – fabrizzio_gz Aug 16 '20 at 20:40
  • \$\begingroup\$ @superb rain Link to this kata is: codewars.com/kata/5f134651bc9687000f8022c4/python \$\endgroup\$ – sid_mallick Aug 17 '20 at 6:42
  • \$\begingroup\$ @fabrizzio_gz, I have the same feeling... looking for it. \$\endgroup\$ – sid_mallick Aug 17 '20 at 6:44
  • 1
    \$\begingroup\$ @fabrizzio_gz The problem statement ends with "tip: you can solve this using smart brute-force". And I just did. Now I can finally look at the code here and perhaps review :-). At first sight it looks similar to mine. \$\endgroup\$ – superb rain Aug 17 '20 at 16:57
  • \$\begingroup\$ The challenge is not really understandable with the link, I think you omitted the most important part in the construction of the sequence. \$\endgroup\$ – Helena Aug 18 '20 at 6:23
2
\$\begingroup\$

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 <= 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 >= 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 > 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 <= 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 >= n:
            x = (arr_sum - n) // i
            return arr_len + arr[i] - (x+1)
        arr_len += arr[i]
        if arr_len < 70_000:
            arr += [i] * arr[i]

And... you can combine the above two improvements, i.e., apply that if arr_len < 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.

\$\endgroup\$
1
  • \$\begingroup\$ Thank you so much for your brilliant optimization suggestions. I am a newbie in the coding world and learnt quite a lot from your solutions. \$\endgroup\$ – sid_mallick Aug 19 '20 at 6:16

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.