Ката Codewars - «Повторяющаяся последовательность»

Aug 17 2020

Я пытаюсь решить следующую ката Codewars .

Нам дан список как
seq = [0, 1, 2, 2]

Нам нужно будет написать функцию, которая, во-первых, будет добавлять элементы в список, используя следующую логику.
если n = 3 и как seq [n] = 2, новый список будет иметь вид seq = [0, 1, 2, 2, 3, 3],
если n = 4 и как seq [n] = 3, новый список будет будет seq = [0, 1, 2, 2, 3, 3, 4, 4, 4],
если n = 5 и как seq [n] = 3, новый список будет seq = [0, 1, 2, 2 , 3, 3, 4, 4, 4, 5, 5, 5] и так далее.

Затем функция вернет n-й элемент списка.

Некоторые элементы списка:
[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]

Ограничение для Python:
0 <= n <= 2 ^ 41

Мой код успешно работает в моей системе для любого значения n, включая n = 2 ^ 41 (в пределах 2,2 с). Но это время истекает в Codewars. Может ли кто-нибудь помочь мне в оптимизации моего кода? Заранее спасибо.

Мой код:

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]

Ответы

2 superbrain Aug 18 2020 at 00:13

Составление списка занимает много времени. Обратите внимание, что вы arr += [i] * arr[i]повторяете одно iи то же значение снова и снова. Его можно сильно сжать, просто сохранив repeat(i, arr[i])вместо этого. Это выполняется примерно за 6 секунд, что значительно меньше их 12-секундного лимита:

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

Обратите внимание, что в этом n > 3случае мы начинаем уже с числа 3, добавляя его дважды в список. Таким образом, [0, 1, 2, 2]нам нужно только начало последовательности [2], поэтому я начинаю с arr = [[2]](что короче [repeat(2, 1)], и chainэто не имеет значения).


В качестве альтернативы ... обратите внимание, что вы растягиваетесь arrнамного быстрее, чем потребляете. Для n = 2 41 вы увеличиваете его до более чем 51 миллиона элементов, но на самом деле вы читаете меньше первых 70 тысяч. Так что на этом этапе вы можете перестать по-настоящему расширять список. Это выполняется примерно за 4,7 секунды:

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]

И ... вы можете объединить два вышеупомянутых улучшения, то есть применить это if arr_len < 70_000:к repeat-version. Затем это происходит примерно за 4,5 секунды.

Результаты тестов на моем ПК для n = 2 41 :

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)

О, и комментарий стиля: вы делаете это дважды:

if ...:
    return ...
else:
    ...

elseИ отступы всего остального кода не нужны и я бы избежать. Я сделал это в приведенных выше решениях.