2 files changed, 20 insertions, 6 deletions
diff --git a/yt/yt/core/misc/heap-inl.h b/yt/yt/core/misc/heap-inl.h
index c26cf45d53a..da1bf99fc0a 100644
--- a/yt/yt/core/misc/heap-inl.h
+++ b/yt/yt/core/misc/heap-inl.h
@@ -155,26 +155,38 @@ void AdjustHeapBack(TIterator begin, TIterator end)
 }
 
 template <class TIterator, class TComparer, class TOnAssign>
-void ExtractHeap(TIterator begin, TIterator end, TComparer comparer, TOnAssign onAssign)
+void ExtractHeap(TIterator begin, TIterator end, TIterator current, TComparer comparer, TOnAssign onAssign)
 {
     YT_ASSERT(begin != end);
     auto newEnd = end - 1;
-    std::swap(*begin, *newEnd);
-    onAssign(0);
+    if (current == newEnd) {
+        return;
+    }
+    std::swap(*current, *newEnd);
+    onAssign(std::distance(begin, current));
     onAssign(std::distance(begin, newEnd));
-    SiftDown(begin, newEnd, begin, comparer, onAssign);
+    SiftDown(std::move(begin), std::move(newEnd), std::move(current), comparer, onAssign);
+}
+
+template <class TIterator, class TComparer, class TOnAssign>
+void ExtractHeap(TIterator begin, TIterator end, TComparer comparer, TOnAssign onAssign)
+{
+    auto current = begin;
+    ExtractHeap(std::move(begin), std::move(end), std::move(current), comparer, onAssign);
 }
 
 template <class TIterator, class TComparer>
 void ExtractHeap(TIterator begin, TIterator end, TComparer comparer)
 {
-    ExtractHeap(std::move(begin), std::move(end), comparer, [] (size_t) {});
+    auto current = begin;
+    ExtractHeap(std::move(begin), std::move(end), std::move(current), comparer, [] (size_t) {});
 }
 
 template <class TIterator>
 void ExtractHeap(TIterator begin, TIterator end)
 {
-    ExtractHeap(std::move(begin), std::move(end), std::less<>(), [] (size_t) {});
+    auto current = begin;
+    ExtractHeap(std::move(begin), std::move(end), std::move(current), std::less<>(), [] (size_t) {});
 }
 
 template <class TIterator, class TComparer, class TOnAssign>
diff --git a/yt/yt/core/misc/heap.h b/yt/yt/core/misc/heap.h
index 6d65e38cbde..f11c1b4e168 100644
--- a/yt/yt/core/misc/heap.h
+++ b/yt/yt/core/misc/heap.h
@@ -31,6 +31,8 @@ void AdjustHeapBack(TIterator begin, TIterator end);
 //! Extracts the front from the heap on |[begin, end)| by moving
 //! its back to the front and then pushing it down if needed.
 template <class TIterator, class TComparer, class TOnAssign>
+void ExtractHeap(TIterator begin, TIterator end, TIterator current, TComparer comparer, TOnAssign onAssign);
+template <class TIterator, class TComparer, class TOnAssign>
 void ExtractHeap(TIterator begin, TIterator end, TComparer comparer, TOnAssign onAssign);
 template <class TIterator, class TComparer>
 void ExtractHeap(TIterator begin, TIterator end, TComparer comparer);