What is the optimal way to retrieve a list of sorted dates in a list of sorted periods in python?

Question

Suppose I have a list with N sorted dates and M non-overlapping sorted periods with start date, end date (inclusive), and a tax rate for example. I have to make an efficient algorithm retrieve all tax rates for all dates. If there is not period including this date it should raise an error.

Given the brute-force approach I could have a O(N * M) with two nested loops. It is possible to break the inner loop when one date is found (maintains the code worst-case complexity). Another optimization would be to store the index of last period, since the lists are sorted, then I believe I got O(N + M). Is there a more optimal way of doing that? Maybe using other data structures?

Working code in Python:

import collections
import datetime
import sys
from typing import List

RatePeriod = collections.namedtuple("RatePeriod", ["start_date", "end_date", "rate"])

periods = [
    RatePeriod(datetime.datetime(2019, 1, 3), datetime.datetime(2019, 4, 1), 10.7),
    RatePeriod(datetime.datetime(2019, 4, 2), datetime.datetime(2019, 12, 2), 20.5),
    RatePeriod(datetime.datetime(2019, 12, 3), datetime.datetime(2020, 1, 2), 37.8),
]

def get_rates(dates: List[datetime.datetime]) -> List[float]:
    rates = []
    last_period = 0
    for idx, date in enumerate(dates, 1):
        for idx2 in range(last_period, len(periods)):
            period = periods[idx2]
            last_period = idx2
            if period.start_date <= date <= period.end_date:
                rates.append(period.rate)
                break
        if len(rates) < idx:
            sys.exit("No period found for date: {}".format(date))
    return rates


series = [
    datetime.datetime(2019, 2, 20),
    datetime.datetime(2019, 3, 6),
    datetime.datetime(2019, 12, 14),
]

result = get_rates(series)
expected = [10.7, 10.7, 37.8]
assert result == expected

Answer 1

Your algorithm seems fine, it seems you brought it down in complexity enough. I have thought about it and couldn't think of anything better.

Still, the code can be rewritten to be more Pythonic:

import collections
import sys
import timeit
from datetime import datetime as dt
from typing import List

RatePeriod = collections.namedtuple("RatePeriod", ["start_date", "end_date", "rate"])

periods = [
    RatePeriod(dt(2019, 1, 3), dt(2019, 4, 1), 10.7),
    RatePeriod(dt(2019, 4, 2), dt(2019, 12, 2), 20.5),
    RatePeriod(dt(2019, 12, 3), dt(2020, 1, 2), 37.8),
    RatePeriod(dt(2020, 1, 3), dt(2020, 12, 2), 41.3),
    RatePeriod(dt(2020, 12, 3), dt(2021, 1, 2), 52.7),
]


series = [
    dt(2019, 2, 20),
    dt(2019, 3, 6),
    dt(2020, 1, 5),
    dt(2020, 12, 5),
    # dt(2022, 1, 1),  # error, no period found
]


def get_rates(dates: List[dt]) -> List[float]:
    rates = []
    last_period = 0
    for idx, date in enumerate(dates, 1):
        for idx2 in range(last_period, len(periods)):
            period = periods[idx2]
            last_period = idx2
            if period.start_date <= date <= period.end_date:
                rates.append(period.rate)
                break
        if len(rates) < idx:
            sys.exit("No period found for date: {}".format(date))
    return rates


def get_rates_generator(dates: List[dt]) -> List[float]:
    last_period = 0
    for date in dates:
        for idx_period, period in enumerate(periods[last_period:], start=last_period):
            if period.start_date <= date <= period.end_date:
                last_period = idx_period
                break
        else:
            sys.exit(f"No period found for date: {date}")
        yield period.rate


result = get_rates(series)
result_generator = list(get_rates_generator(series))

setup = "from __main__ import get_rates, get_rates_generator, series"
print("Old:", timeit.timeit("get_rates(series)", setup=setup))
print("New:", timeit.timeit("list(get_rates_generator(series))", setup=setup))

expected = [10.7, 10.7, 41.3, 52.7]

assert result == result_generator == expected

where the printed output will be somewhere in the ballpark of

Old: 2.3620867
New: 2.3765742000000003

aka, the rewritten function is not actually faster. A couple notes on the new, suggested approach:

• imported datetime.datetime as dt for readability; from datetime import datetime is likely confusing and should be avoided
• handling the case of not finding an item is nicely handled using for/else. In fact, it is one of the prime uses of for/else. The else block is run if no break was encountered for the entire loop. Since there is already a break in place, it is straightforward to implement the else block. This also gets rid of having to enumerate over dates.
• instead of building and returning a list, a generator only yields on request. Having found a match, the loop is broken out of and yield is hit. The function exits and returns to its saved state on the next iteration. As such, list() can be called on the generator object to exhaust it fully and receive a list, as in your code. If no full list is desired, a generator is lighter than a full list.
• the enumerate function, together with list slicing, can do what your range code did. This way, the new code is a bit dense on that line; since it is also verbose, I think it's manageable. I find it to be more readable.