class Heap:
def __init__(self, data=None, min_heap=True):
self._min = min_heap
self._data = list(data) if data is not None else []
for i in range((len(self._data) - 2) // 2, -1, -1):
self._sift_down(i)
def __len__(self):
return len(self._data)
def peek(self):
if not self._data:
raise IndexError("peek from empty heap")
return self._data[0]
def push(self, val):
self._data.append(val)
self._sift_up(len(self._data) - 1)
def pop(self):
if not self._data:
raise IndexError("pop from empty heap")
self._swap(0, len(self._data) - 1)
val = self._data.pop()
self._sift_down(0)
return val
def pushpop(self, val):
if self._data and self._cmp(self._data[0], val):
val, self._data[0] = self._data[0], val
self._sift_down(0)
return val
def _cmp(self, a, b):
return a < b if self._min else a > b
def _parent(self, i):
return (i - 1) // 2
def _left(self, i):
return 2 * i + 1
def _right(self, i):
return 2 * i + 2
def _swap(self, i, j):
self._data[i], self._data[j] = self._data[j], self._data[i]
def _sift_up(self, i):
while i > 0:
p = self._parent(i)
if self._cmp(self._data[i], self._data[p]):
self._swap(i, p)
i = p
else:
break
def _sift_down(self, i):
n = len(self._data)
while True:
best = i
l, r = self._left(i), self._right(i)
if l < n and self._cmp(self._data[l], self._data[best]):
best = l
if r < n and self._cmp(self._data[r], self._data[best]):
best = r
if best == i:
break
self._swap(i, best)
i = best
n, k = [int(x) for x in input().split()]
heights = [int(x) for x in input().split()]
visited = set()
heap = Heap([(x, i) for i, x in enumerate(heights)], False)
cost = 0
while len(heap) > 0:
val, i = heap.pop()
if i in visited:
continue
visited.add(i)
if i > 0 and heights[i - 1] + k < val:
d = val - heights[i - 1] - k
cost += d
heights[i - 1] += d
heap.push((heights[i - 1], i - 1))
if i < len(heights) - 1 and heights[i + 1] + k < val:
d = val - heights[i + 1] - k
cost += d
heights[i + 1] += d
heap.push((heights[i + 1], i + 1))
print(cost)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | class Heap: def __init__(self, data=None, min_heap=True): self._min = min_heap self._data = list(data) if data is not None else [] for i in range((len(self._data) - 2) // 2, -1, -1): self._sift_down(i) def __len__(self): return len(self._data) def peek(self): if not self._data: raise IndexError("peek from empty heap") return self._data[0] def push(self, val): self._data.append(val) self._sift_up(len(self._data) - 1) def pop(self): if not self._data: raise IndexError("pop from empty heap") self._swap(0, len(self._data) - 1) val = self._data.pop() self._sift_down(0) return val def pushpop(self, val): if self._data and self._cmp(self._data[0], val): val, self._data[0] = self._data[0], val self._sift_down(0) return val def _cmp(self, a, b): return a < b if self._min else a > b def _parent(self, i): return (i - 1) // 2 def _left(self, i): return 2 * i + 1 def _right(self, i): return 2 * i + 2 def _swap(self, i, j): self._data[i], self._data[j] = self._data[j], self._data[i] def _sift_up(self, i): while i > 0: p = self._parent(i) if self._cmp(self._data[i], self._data[p]): self._swap(i, p) i = p else: break def _sift_down(self, i): n = len(self._data) while True: best = i l, r = self._left(i), self._right(i) if l < n and self._cmp(self._data[l], self._data[best]): best = l if r < n and self._cmp(self._data[r], self._data[best]): best = r if best == i: break self._swap(i, best) i = best n, k = [int(x) for x in input().split()] heights = [int(x) for x in input().split()] visited = set() heap = Heap([(x, i) for i, x in enumerate(heights)], False) cost = 0 while len(heap) > 0: val, i = heap.pop() if i in visited: continue visited.add(i) if i > 0 and heights[i - 1] + k < val: d = val - heights[i - 1] - k cost += d heights[i - 1] += d heap.push((heights[i - 1], i - 1)) if i < len(heights) - 1 and heights[i + 1] + k < val: d = val - heights[i + 1] - k cost += d heights[i + 1] += d heap.push((heights[i + 1], i + 1)) print(cost) |
English