Разбиение на интервалы дат

Пару дней назад участвовал в квалификации Yandex Cup 2020 (по ее итогам все-таки пробился в финал). Первой из шести задач была задача "Разбиение на интервалы дат". Вот ее условие:

Миша работает в команде Яндекс.Маркета, которая предоставляет производителям товаров аналитику о продажах. Сейчас Миша разбирается с периодизацией: нужно собирать данные по дням, неделям, месяцам, кварталам и годам. От клиентов приходят запросы, в которых указан период детализации и интервал: начальная и конечная даты. Так что первоначально Мише нужно разбить интервал на периоды. Так, если клиент хочет данные с 2020-01-10 по 2020-03-25 с детализацией по месяцам, то ему вернутся данные за три периода: c 2020-01-10 по 2020-01-31, с 2020-02-01 по 2020-02-29 и с 2020-03-01 по 2020-03-25. Помогите Мише, а то ему еще диплом писать надо!

Всего нужно поддержать пять видов временных интервалов:

  • WEEK — неделя с понедельника по воскресенье.
  • MONTH — месяц.
  • QUARTER — интервалы в три месяца: январь — март, апрель — июнь, июль — сентябрь, октябрь — декабрь.
  • YEAR — год c 1 января по 31 декабря.
  • LAST_SUNDAY_OF_YEAR — интервал, начинающийся в последнее воскресенье декабря.

Формат ввода

В первой строке дан типа интервала type — строка, принимающая одно из следующих значений: WEEK, MONTH, QUARTER, YEAR, LAST_SUNDAY_OF_YEAR. Во второй строке через пробел даны начальная и конечная даты start и end (start ≤ end ) в формате yyyy-MM-dd. Гарантируется, что обе даты лежат в промежутке с 1 января 2000 года по 31 декабря 3999 года включительно.

Формат вывода

В первой строке ответа выведите одно целое число N — количество промежутков. В последующих N строках на i-й строке выведите через пробел дату начала и конца i-го промежутка в формате yyyy-MM-dd. Промежутки должны выводиться в порядке возрастания начальной даты.

Пример 1

Ввод:

MONTH
2020-01-10 2020-03-25

Вывод:

3
2020-01-10 2020-01-31
2020-02-01 2020-02-29
2020-03-01 2020-03-25

Пример 2

Ввод:

WEEK
2020-01-26 2020-03-23

Вывод:

10
2020-01-26 2020-01-26
2020-01-27 2020-02-02
2020-02-03 2020-02-09
2020-02-10 2020-02-16
2020-02-17 2020-02-23
2020-02-24 2020-03-01
2020-03-02 2020-03-08
2020-03-09 2020-03-15
2020-03-16 2020-03-22
2020-03-23 2020-03-23

Пример 3

Ввод:

LAST_SUNDAY_OF_YEAR
2018-01-01 2021-01-01

Вывод:

4
2018-01-01 2018-12-29
2018-12-30 2019-12-28
2019-12-29 2020-12-26
2020-12-27 2021-01-01

А теперь перейдем к коду. Написал его с использованием стандартных модулей, прогнал на стандартных тестах и... получил Wrong Answer на 8 (скрытом) тесте. Тесты так и не открыли, разбора задач нету, а ошибку найти хочется.

from datetime import datetime, timedelta
import calendar

result = []
interval = input()
start, end = [datetime.strptime(date, '%Y-%m-%d') for date in input().split()]
if start == end:
    print(1)
    print(start.strftime('%Y-%m-%d') + ' ' + end.strftime('%Y-%m-%d'))
else:
    while True:
        after = start
        if interval == 'MONTH':
            after = after.replace(
                day=list(calendar.monthrange(start.year, start.month))[1])
        elif interval == 'WEEK':
            while after.isoweekday() != 7:
                after += timedelta(days=1)
        elif interval == 'YEAR':
            after = after.replace(day=31, month=12)
        elif interval == 'QUARTER':
            while not ((after.day == 31 and after.month == 3)
                       or (after.day == 30 and after.month == 6)
                       or (after.day == 30 and after.month == 9)
                       or (after.day == 31 and after.month == 12)):
                after += timedelta(days=1)
        elif interval == 'LAST_SUNDAY_OF_YEAR':
            after = after.replace(day=31, month=12)
            while after.isoweekday() != 6:
                after -= timedelta(days=1)
            if after < start:
                after = after.replace(day=31, month=12, year=after.year + 1)
                while after.isoweekday() != 6:
                    after -= timedelta(days=1)
        if after < end:
            result.append(start.strftime('%Y-%m-%d') +
                          ' ' + after.strftime('%Y-%m-%d'))
            start = after
            start += timedelta(days=1)
        else:
            result.append(start.strftime('%Y-%m-%d') +
                          ' ' + end.strftime('%Y-%m-%d'))
            break

    print(len(result))
    print('\n'.join(result))

Надеюсь на вашу помощь :)


Ответы (0 шт):