import heapq
class Solution:
def supplyWagon(self, supplies: List[int]) -> List[int]:
# m = {}
# for i in range(len(supplies)):
# m[i] = supplies[i]
# def findNext(m, i, l):
# r = i + 1
# while r < l and r not in m:
# r += 1
# return r
# l = len(supplies)
# bound = l
# while l != (len(supplies) >> 1):
# i1 = 0
# i2 = 0
# r1 = 0
# r2 = 0
# ma = 100000
# while i1 != bound - 1:
# i2 = findNext(m, i1, bound)
# if m[i1] + m[i2] < ma:
# ma = m[i1] + m[i2]
# r1 = i1
# r2 = i2
# i1 = i2
# m[r1] += m[r2]
# del m[r2]
# if r2 == bound - 1:
# bound = r1 + 1
# l -= 1
# arr = []
# for i, j in m.items():
# arr.append([i,j])
# arr.sort(key=lambda x:x[0])
# r = [x[1] for x in arr]
# return r
def findPrev(supplies, i):
i -= 1
while i >= 0 and supplies[i] == -1:
i -= 1
return i
def findNext(supplies, i):
i += 1
while i < len(supplies) and supplies[i] == -1:
i += 1
return i
arr = [[supplies[i] + supplies[i + 1], i, i + 1] for i in range(len(supplies) - 1)]
heapq.heapify(arr)
l = (len(supplies) + 1) >> 1
i = 0
while i < l:
t = heapq.heappop(arr)
while supplies[t[1]] + supplies[t[2]] != t[0]:
t = heapq.heappop(arr)
supplies[t[1]] += supplies[t[2]]
supplies[t[2]] = -1
prev = findPrev(supplies, t[1])
next = findNext(supplies, t[2])
if prev != -1:
heapq.heappush(arr, [supplies[prev] + supplies[t[1]], prev, t[1]])
if next != len(supplies):
heapq.heappush(arr, [supplies[next] + supplies[t[1]], t[1], next])
i += 1
r = []
for s in supplies:
if s != -1:
r.append(s)
return r
import heapq
class Solution:
def supplyWagon(self, supplies: List[int]) -> List[int]:
# Helper function to find the previous non-merged wagon
def findPrev(supplies, i):
i -= 1
while i >= 0 and supplies[i] == -1:
i -= 1
return i
# Helper function to find the next non-merged wagon
def findNext(supplies, i):
i += 1
while i < len(supplies) and supplies[i] == -1:
i += 1
return i
# Initialize the min-heap with the sum of every two adjacent wagons
arr = [[supplies[i] + supplies[i + 1], i, i + 1] for i in range(len(supplies) - 1)]
heapq.heapify(arr)
# Calculate the target length after all merges
l = (len(supplies) + 1) >> 1
i = 0
while i < l:
# Get the smallest pair from the heap
t = heapq.heappop(arr)
# Check if the pair is still valid
while supplies[t[1]] + supplies[t[2]] != t[0]:
t = heapq.heappop(arr)
# Merge the two wagons
supplies[t[1]] += supplies[t[2]]
supplies[t[2]] = -1
# Update the heap with new possible pairs
prev = findPrev(supplies, t[1])
next = findNext(supplies, t[2])
if prev != -1:
heapq.heappush(arr, [supplies[prev] + supplies[t[1]], prev, t[1]])
if next != len(supplies):
heapq.heappush(arr, [supplies[next] + supplies[t[1]], t[1], next])
i += 1
# Collect non-merged wagons for the result
r = []
for s in supplies:
if s != -1:
r.append(s)
return r