Apply additions on a base array based on date ranges

Question

please roast my code, I hope this is the good place to look for some advice on where I can improve.

Problem definition:

Apply modifications on values depending on date ranges.

Data:

Base array - holds dictionaries in the following format:

[{'date': 20200101, 'value': 1}, {'date': 20200102, 'value': 2}]

Modifier array - holds dictionaries with similar format:

[{'date': 20200101, 'value': 1}, {'date': 20200201, 'value': 2}]

Edit: Both arrays are holding values with dates sorted in ascending order.

Goal: Add the respective value of the modifier array to the base array lining up the date ranges. Dates are exclusive, for example when the modifier array contains 2020-01-01 you have to add the value '1' to all values in the base array that have a date less than 2020-01-01. Base array has a lot of elements while modifier array relatively few. In practice this splits the base array into a couple of date ranges. If the last modification date is less than the date in base array no modification is required.

My solution:

This is assuming the comparing of dates will work, I have translated this from perl for an easier read.

mod_index = 0 
mod_size = len(mod_arr)

for elem in base_arr:
  if elem['date'] > mod_arr[mod_size - 1]['date']:
    break
  else:
    if elem['date'] < mod_arr[mod_index]['date']:
      elem['value'] += mod_arr[mod_index]['value']
    else:
      elem['value'] += mod_arr[mod_index + 1]['value']
      mod_index += 1

Answer 1

To preface, PEP8 recommends 4 spaces instead of 2 spaces for indentation, so I've written all below code as such.

There is a subtle bug in your program (and may be nonexistent at all, the requirements are ambiguous here), but:

if elem['date'] < mod_arr[mod_index]['date']:
    elem['value'] += mod_arr[mod_index]['value']
else:
    elem['value'] += mod_arr[mod_index + 1]['value']
    mod_index += 1

assumes that mod_array[mod_index + 1]["date"] > mod_array[mod_index]["date"], though your requirements never say that "date" is strictly increasing (and mod_array being sorted doesn't imply this).

We can handle this by changing the if / else to a while:

while elem['date'] >= mod_arr[mod_index]['date']:
    mod_index += 1

elem['value'] += mod_arr[mod_index]['value']

Note that this also increments mod_index before the mutation of elem["value"], allowing us to de-duplicate its mutation.

Since we are only moving forward in mod_arr, we actually don't need mod_index at all. We can do mod_array_iter = iter(mod_array), and replace mod_index += 1 with mod = next(mod_array_iter). This also helps us avoid any off-by-one errors.

Another benefit of using iter here is it opens up an easy way to exit the function if mod_array is empty (implying no mutations to base_arr should be done, the original program didn't account for this) or if we've traversed through all of mod_array_iter (implying necessary mutations to base_arr have been done, original program accounted for this with the if and inner break). We can use a try / except for both of those cases, and get rid of the initial if with inner break.

Final code ends up looking like (wrapped is a function for reusability):

def mutate_base_arr(base_arr, mod_arr):
    mod_arr_iter = iter(mod_arr)

    try:
        mod = next(mod_arr_iter)

        for elem in base_arr:
            while elem['date'] >= mod['date']:
                mod = next(mod_arr_iter)

            elem['value'] += mod['value']
    except StopIteration:
        return

If you really wanted to take this a step further, you could get rid of the while and duplicate nexts by using itertools, yielding this version of the function:

def mutate_base_arr(base_arr, mod_arr):
    mod_arr_iter = iter(mod_arr)

    try:
        for elem in base_arr:
            valid_mod = next(mod for mod in mod_arr_iter if elem["date"] < mod["date"])
            mod_arr_iter = itertools.chain([valid_mod], mod_arr_iter)

            elem['value'] += valid_mod['value']
    except StopIteration:
        return

Though I'm not sure if I'd recommend this, as having to prepend valid_mod with mod_arr_iter each iteration so it can be re-checked next iteration of the for is a bit unusual.