How to optimize this algorithm for repeatedly finding and updating the minimum of an array?

Viewed 99

Input description

The first line of my input contains two numbers: and , denoting respectively the number of containers (2 ≤ ≤ 20,000) and the number of deliveries (not more than 100,000)

In the following lines there are pairs of numbers: and , where is the number of wagons in the delivery (no more than 20,000) and is the number of gravel in each wagon (no more than 20,000).

The total number of wagons from all deliveries will not exceed 1,000,000.
There will be no situation in any of the inputs where the container will contain more than 1,000,000,000 gravel.

Problem statement

Wagons are emptied one at a time, evenly into two consecutive containers. These are the rules for determining which two adjacent containers are chosen (where containers[] denotes a list of the n current container sizes):

  1. Among all pairs (containers[i], containers[i+1]) with 0 <= i < n - 1, choose the pair which minimizes min(containers[i], containers[i+1])
  2. If there is a tie, choose among the tied pairs, the pair which minimizes max(containers[i], containers[i+1]) (i.e. minimizes the second smallest element in the pair)
  3. If there is still a tie, choose the pair (containers[i], containers[i+1]) with minimum index i among the tied pairs.

The gravel is distributed equally into the pair. If there is an odd number of gravel in a wagon, the container with the smaller ID gains the surplus (1 pebble).

At first, all containers are empty.

For example:

lista = [0, 4, 3, 2, 0, 5, 4, 0, 3, 4, 2, 0]

My first smallest value in that list is 0 so there are 6 possibilities for different pairs

0,4 with indices 0 and 1  
0,2 with indices 3 and 4  
0,5 with indices 4 and 5  
4,0 with indices 6 and 7  
2,0 with indices 10 and 11  

My second smallest value in this 6 possibilities is 2 so I have two options

0,2 with indices 3 and 4  
2,0 with indices 10 and 11  

Indexes 3 and 4 are less than 10 and 11 so we will be dropping the wagons into containers with indexes [3, 4]

In the output:
Each of the lines of output data should contain a single number - the number of pebbles in a container. The order of listing should be the same as the order of the identifiers.

Example
For the example input shown below:

6 4  
4 3  
1 2  
2 14  
4 3 

the correct answer is:

6  
10  
12  
7  
11  
8  

Explanation of the example

  1. In the first delivery:

    • The first wagon can be unloaded into pairs of containers (1, 2), (2, 3), (3, 4), (4, 5), (5, 6).
      We choose the first one and the values ​​of the filling are (2 1 0 0 0 0).
    • Second wagon can be put into pairs of containers (3, 4), (4, 5), (5, 6).
      We choose the first one and the filling values ​​are (2 1 2 1 0 0).
    • The third wagon can be put into a pair of containers (5, 6), which will give us the values ​​of the fills (2 1 2 1 2 1).
    • The fourth wagon may be discharged into par containers (1, 2), (2, 3), (3, 4), (4, 5), (5, 6). We choose the first one and the filling values is (4 2 2 1 2 1).
  2. The second delivery changes the capacity of the containers as follows: (4 2 3 2 2 1).

  3. In the third delivery:

    • The first wagon may be unloaded only into a pair of containers (5, 6), and then the values ​​for the fills are (4 2 3 2 9 8).
    • The second car can be unloaded into pairs (2, 3) or (3, 4). We choose the first one and the fill values ​​are (4 9 10 2 9 8).
  4. The last delivery changes the fills as follows:

    • (4 9 10 4 10 8),
    • (6 10 10 4 10 8) ,
    • (6 10 12 5 10 8) ,
    • (6 10 12 7 11 8)

Below is my program code (it works, however, it is not optimal for large numbers)

n, d = map(int, input().split())

containers = []
for _ in range(n):
    containers.append(0)

for _ in range(d):
    wagons, gravel = map(int, input().split()) # w, z
    half_gravel = int(gravel/2)
    greater_half = int(gravel/2)+ gravel%2
    for _ in range(wagons):
    
        i = min(range(len(containers) - 1),
        key=lambda i: sorted(containers[i : i+2]))

        containers[i] += greater_half
        containers[i+1] += half_gravel

for i in containers:
    print(i)

Thanks in advance for your help in optimizing the program

2 Answers

You can solve this much more efficiently by using a min-heap (i.e. priority queue), such as the one provided in Python's standard library. The main purpose of a heap is for efficiently tracking the minimum element of a collection, removing that minimum, and adding new elements, which accurately reflects your problem.

Currently, the line:

i = min(range(len(containers) - 1),
        key=lambda i: sorted(containers[i : i+2]))

In your innermost loop is an O(n) cost operation, traversing all containers each time. You can improve this to O(log n) cost by keeping a min-heap of the 3-tuples which is equivalent to:

min_heap = (min(containers[i:i+2]), max(containers[i:i+2]), i)
            for i in range(n-1)

The minimum of these 3-tuples, using the standard lexicographic ordering, gives the next container to receive gravel. Importantly, we don't need to actually rebuild this heap in the loop, just update it

There's one subtlety here: a container can be a member of two of these 3-tuples at a time, so if we update a container pair and push its 3-tuple back onto the heap, there will be an out-of-date 3-tuple in the min-heap with a smaller value than it should have. This is fine: we can either implement a decrease-key function (fairly difficult using this heap), or, just use lazy updating.

For 'lazy updates', when popping from the heap, we check whether the 'reported minimum' of the heap is out of date, remove it, and insert an updated version back in. Using amortized analysis, since this extra heap operation cost occurs cumulatively at most once per wagon, the runtime of this algorithm is still
O(n + total_wagons * log(n)).

Python code:

n, d = map(int, input().split())
containers = [0]*n

# Heap elements will be: (min(pair), max(pair), i) for pair = container[i:i+2]
min_heap = [(0, 0, i) for i in range(n-1)]

for _ in range(d):
    wagons, gravel = map(int, input().split())  # w, z
    half_gravel = gravel // 2
    greater_half = (gravel // 2) + gravel % 2
    for _ in range(wagons):

        # Find minimum element
        while True:
            smaller, larger, index = heapq.heappop(min_heap)
            # Check if entry is out of date
            actual_pair = sorted([containers[index], containers[index+1]])
            if actual_pair != [smaller, larger]:
                heapq.heappush(min_heap, (actual_pair[0], actual_pair[1], index))
            else:
                break

        containers[index] += greater_half
        containers[index + 1] += half_gravel

        new_smaller, new_larger = sorted([containers[index], containers[index+1]])
        # This may cause index+1 entry to be out of date, which is OK.
        heapq.heappush(min_heap, (new_smaller, new_larger, index))

for i in containers:
    print(i)

Note: This lazy updating strategy is useful in general to avoid needing decrease-key, and is more commonly used in Dijstra's algorithm. However, it only works if elements in your queue may have stored-values smaller than their true values, not the other way around.

The while True: loop looks a bit strange, but it's used because there may be up (in theory) O(n) out-of-date elements in the heap, so we need to keep popping, checking, and updating if the reported minimum is out-of-date.

How can we tell the number of elements in the heap doesn't change between wagons? At the start of every iteration of that loop, there has been no net change in the number of elements. In each iteration, we perform a pop at the start. Then, either we:

  • Find an out-of-date element, perform a push and continue (net 0),
  • Break (net -1 elements).

So after the while True loop, we are always at net -1 elements. After the loop, we always perform a push, bringing us back to no net change in heap size.

I endorse @kcsquared's approach, but will post a variation that's less "clever" in all respects, and so more likely to be obvious ;-) The output matches what you expect in the single test case you gave.

EDIT: as pointed out in a comment, this doesn't always solve the problem as stated. I'm going to let the answer sit here, though, because it's instructive in some other aspects.

from heapq import heapreplace

containers = [0] * n
h = [(0, i) for i in range(n)]

for wagons, gravel in [(4, 3),
                       (1, 2),
                       (2, 14),
                       (4, 3),
                      ]:
    half_gravel = gravel // 2
    greater_half = gravel - half_gravel
    for _ in range(wagons):
        while True:
            val, i = h[0]
            if val != containers[i]:
                heapreplace(h, (containers[i], i))
            else:
                break
        if i == len(containers) - 1:
            i -= 1
        elif i and containers[i-1] <= containers[i+1]:
            i -= 1
        containers[i] += greater_half
        containers[i+1] += half_gravel

for i in containers:
    print(i)
assert containers == [6, 10, 12, 7, 11, 8]

Related