优化 python 代码 - 在未排序的列表中搜索，需要保留嵌套循环作为索引位置

Question

I have a list (stockData) with some random positive integers then another query list (queries) that contains some positions(indexes counted from 1 and not 0) within stockData.

Based on the 'queries', The code must identify a closest value smaller than the one at a give position. If there is a clash/tie, the smaller index position is preferred. If no such value can be identified on either side of position, then -1 should be returned.

Example1

stockData = [5,6,8,4,9,10,8,3,6,4]
queries = [6,5,4]

At position 6 (index=5) in stockData, value is 10, both position 5 (index=4) and position 7 (index=6) are smaller values (9 and 8 resp.) than 10, hence we choose the one at a lesser position -> [5]
At position 5 (index=4) in stockData, value is 9, position 4 (index=3) is smaller hence we choose position 4 -> [5,4]
At position 4 (index=3) in stockData, value is 4, position 8 (index=7) is only value smaller hence we choose position 8 -> [5,4,8]

Output
[5,4,8]

Example2

stockData = [5,6,8,4,9,10,8,3,6,4]
queries = [3,1,8]

Here, at position 8, the value 3 is the smallest of the list, so we return -1

Output
[2,4,-1]

Constraints

1 <= n <= 10^5 (n is the length of stockData)
1 <= stockData[i] <= 10^9 
1 <= q <= 10^5 (q is the length of queries)
1 <= queries[i] <= 10^9 

My code is working fine, but taking too long to execute. Looking for any help with optimizing it or any other corrective measures. Thanks !!

My Code:

def predictAnswer(stockData, queries):
    # convert days to index
    query_idx = [val-1 for val in queries]

    # output_day_set
    out_day = []

    for day in query_idx:
        min_price = stockData[day]
        day_found = False
        
        for i in range(1, max(day,len(stockData)-day)):
            prev = day-i
            nxt = day+i
            if prev >= 0 and stockData[prev] < min_price:
                out_day.append(prev+1)                        
                day_found = True
                break
            if nxt < len(stockData) and stockData[nxt] < min_price:
                out_day.append(nxt+1)                        
                day_found = True
                break
        
        if not day_found:
            out_day.append(-1)
    
    return out_day

Solution Approach as suggested by @Carlos:

for idx, val in enumerate(stockData):
        if val in stockDict.keys():
            stockDict[val].append(idx)
        else:
            stockDict[val] = [idx]
    
    stockDict_sortedKeys = sorted(stockDict)

For the sample input:

stockData = [5,3,4,6,5,4,6,8,9,5,7]

Output:

stockDict = {5: [0], 6: [1, 8], 8: [2, 6], 4: [3, 9], 9: [4], 10: [5], 3: [7]}
stockDict_sortedKeys = [3, 4, 5, 6, 8, 9, 10]

Now may be for a for a given query 'day', I would need to find the corresponding sortedData[day] and find the closest index for all the values smaller?

Answer 1

You can iterate on your stockData to get the most recent previous smaller value, by keeping a heap of local minima:

import heapq

class Data:
    def __init__(self, day, value):
        self.day = day
        self.value = value
    def __lt__(self, other):
        return self.value >= other.value
    def __repr__(self):
        return f'Data(day={self.day}, value={self.value})'

def precalc(days):
    result = []
    heap = []
    for day, value in days:
        data = Data(day, value)
        while heap and heap[0] < data:
            heapq.heappop(heap)
        result.append(heap[0].day if heap else -1)
        heapq.heappush(heap, data)
    return result

The complexity is O(n log(n)) , as each day can only be pushed/popped once.

To get the expected answer, you have to do it in the other direction also:

def compare(f, b, q):
    if f < 0: return b
    if b < 0: return f
    if q-f <= b-q:
        return f
    else:
        return b

def predictAnswer(stockData, queries):
    forward = precalc(enumerate(stockData, 1))
    backward = list(reversed(precalc(reversed(list(enumerate(stockData, 1))))))
    return [compare(forward[q-1], backward[q-1], q) for q in queries]

This is still O(n log(n)) .