Python-оптимизация алгоритма динамического программирования из codechef

Версия 00:04, 1 мая 2022

При разборе ваших решений, столкнулся с одной из попыток, с результатом «вроде все сделано правильно, но система не пропускает, наверно что-то не так с системой или вообще питон не потянет, вот на плюсах все решают и все хорошо».

Попробую кратко показать, как можно системно поработать с улучшением решения, даже не особо вьезжая в алгоритм (ну при условии, что идея будет правильна).

import sys
 
digits = 19
# globals
input_str = ""
dp = [[[-1, -1], [-1, -1]] for i in range(digits)]
N = None
 
 
def get_input():
    lines = sys.stdin.readlines()
    content = " ".join(lines).strip()
    for token in content.split():
        yield token
 
tc_ = get_input()
next(tc_)
 
p10 = [10**i for i in range(digits)]
 
 
def calculate(i, taken: int, ls: int) -> int:
    '''
    i — индекс цифры
    taken — встречали ли простые цифры раньше
    ls — в текущем индексе можно делать полный перебор
    '''
    global input_str
    global dp
    global N
 
    res = 0
    if ls:
        limit = 9
    else:
        limit = int(input_str[i])
 
    for x in range(limit + 1):
        newtaken = taken or (x == 2) or (x == 3) or (x == 5) or (x == 7)
        newls = ls or (x != limit)
        newi = i + 1
 
        if newi == N:
            ret = newtaken
        elif newls and newtaken:
            ret = p10[N - newi]
        else:
            cached = dp[newi][newtaken][newls]
            if cached != -1:
                ret = cached
            else:
                ret = calculate(newi, newtaken, newls)
 
        res += ret
 
    dp[i][taken][ls] = res
    return res
 
 
for input_str in tc_:
    N = len(input_str)
    for j in range(digits):
        dp[j][0][0] = -1
        dp[j][1][0] = -1
        dp[j][1][1] = -1
        dp[j][0][1] = -1
    print(calculate(0, 0, 0))

Видим вроде типичный алгоритм ДП, смотрим описание задачи Codechef/DIGPRIME, переходим по ссылке Editorial на разбор и разъяснение задачи, но даже если этой ссылки не будет, смотрим список принятых решений:

Берем из них любое принятое CPP-решение, сохраняем его в файл digprime-good.cpp компилируем его

 gcc -g -o digprime-good digprime-good.cpp -lstdc++

Итак, у нас есть референсное решение.

Смотрим на описание задачи, особенно на секцию «ограничения»…

и пишем примитивных генератор «digprime-generate.py», такой, чтобы задействовать все ограничения (вдруг все проходит на минимальных данных, но где-то что-то переполняется на самых крайних случаях), плюс большой тестсет даст возможность разумно измерять время работы.

import numpy as np
num = 100000
print(num)
for t in range(num):
    print(np.random.randint(1,1000000000000000000))

Генерим наш тестовый набор:

python digprime-generate.py > big-samples.txt

Генерим результаты нашего алгоритма и референсной реализации на тех же входных данных:

digprime-good < big-samples.txt > reference-good.txt
python digprime.py < big-samples.txt > big-our-results.txt

Сравниваем («meld», «winmerge», «fc», ) — используйте то, что ставится на вашу ось и есть под рукой, но в данном случае, совпадение добайтовое даже по ответу «md5sum»:

md5sum big-our-results.txt
  e602cd2d36e9c6749764cace13774bdc  big-our-results.txt

md5sum reference-good.txt
  e602cd2d36e9c6749764cace13774bdc  reference-good.txt

Вроде же все в абсолютном порядке, но что выдает codechef?

Ошибка «NZEC» — «какая-то ошибка», увы, без малейших подсказок на чем, и что за ошибка. Но по времени работы — доли секунды, предположим, что что-то сразу с вводом.

Тут на самом деле наблюдались разные проблемы в этих (codechef, spoj) тестовых системах. То что-то не так с буферизацией, и какой-то перевод строки становился пробелом, или длинная строка, и input() не вычитывал ее до конца, то большое количество операций input, а там перенаправление на чтение из файла, много IO операций, срабатывают какие-то контейнерные ограничения или просто будет тормозить.

Поэтому, лучше сразу сделать чтение типа