Repository: KATO-Hiro/Somen-Soupy

Issue:
数列を昇順/降順に並べたときに、k番目までの要素の総和を高速に計算
```py
# -*- coding: utf-8 -*-

from collections import defaultdict
from heapq import heappop, heappush


class SumOfTopKth:
    """Sum of the k-th number from the smallest (largest) to the k-th.

    See:
    https://atcoder.jp/contests/abc306/submissions/42339375
    """

    __slots__ = (
        "_summed",
        "_k",
        "_in",
        "_out",
        "_d_in",
        "_d_out",
        "_freq",
        "_ascending_order",
    )

    def __init__(self, k: int, ascending_order=True) -> None:
        self._k = k
        self._summed = 0
        self._in = []
        self._out = []
        self._d_in = []
        self._d_out = []
        self._ascending_order = ascending_order
        self._freq = defaultdict(int)

    def query(self) -> int:
        return self._summed if self._ascending_order else -self._summed

    def add(self, x: int) -> None:
        if not self._ascending_order:
            x = -x

        self._freq[x] += 1
        heappush(self._in, -x)
        self._summed += x
        self._modify()

    def discard(self, x: int) -> None:
        if not self._ascending_order:
            x = -x
        if self._freq[x] == 0:
            return

        self._freq[x] -= 1

        if self._in and -self._in[0] == x:
            self._summed -= x
            heappop(self._in)
        elif self._in and -self._in[0] > x:
            self._summed -= x
            heappush(self._d_in, -x)
        else:
            heappush(self._d_out, x)

        self._modify()

    def set_k(self, k: int) -> None:
        self._k = k
        self._modify()

    def get_k(self) -> int:
        return self._k

    def _modify(self) -> None:
        while self._out and (len(self._in) - len(self._d_in) < self._k):
            p = heappop(self._out)

            if self._d_out and p == self._d_out[0]:
                heappop(self._d_out)
            else:
                self._summed += p
                heappush(self._in, -p)

        while len(self._in) - len(self._d_in) > self._k:
            p = -heappop(self._in)

            if self._d_in and p == -self._d_in[0]:
                heappop(self._d_in)
            else:
                self._summed -= p
                heappush(self._out, p)

        while self._d_in and self._in[0] == self._d_in[0]:
            heappop(self._in)
            heappop(self._d_in)

    def __len__(self) -> int:
        return len(self._in) + len(self._out) - len(self._d_in) - len(self._d_out)

    def __contains__(self, x: int) -> bool:
        if not self._ascending_order:
            x = -x
        return self._freq[x] > 0


def main():
    import sys

    input = sys.stdin.readline

    n, k, q = map(int, input().split())
    a = [0] * n
    s = SumOfTopKth(k, ascending_order=False)

    for _ in range(q):
        xi, yi = map(int, input().split())
        xi -= 1

        s.discard(a[xi])
        s.add(yi)
        print(s.query())
        a[xi] = yi


if __name__ == "__main__":
    main()

```

Comments:

Patch:
diff --git a/snippets/data_structure/sum_of_top_kth.py b/snippets/data_structure/sum_of_top_kth.py
new file mode 100644
index 0000000..1149736
--- /dev/null
+++ b/snippets/data_structure/sum_of_top_kth.py
@@ -0,0 +1,123 @@
+# -*- coding: utf-8 -*-
+
+"""
+Usage:
+
+#     n, k, q = map(int, input().split())
+#     a = [0] * n
+
+#     # smallest
+#     s = SumOfTopKth(k, ascending_order=True)
+#     # largest 
+#     s = SumOfTopKth(k, ascending_order=False)
+
+#     for _ in range(q):
+#         xi, yi = map(int, input().split())
+#         xi -= 1
+
+#         s.discard(a[xi])
+#         s.add(yi)
+#         print(s.query())
+#         a[xi] = yi
+"""
+
+from collections import defaultdict
+from heapq import heappop, heappush
+
+
+class SumOfTopKth:
+    """Sum of the k-th number from the smallest (largest) to the k-th.
+
+    See:
+    https://atcoder.jp/contests/abc306/submissions/42339375
+    """
+
+    __slots__ = (
+        "_summed",
+        "_k",
+        "_in",
+        "_out",
+        "_d_in",
+        "_d_out",
+        "_freq",
+        "_ascending_order",
+    )
+
+    def __init__(self, k: int, ascending_order=True) -> None:
+        self._k = k
+        self._summed = 0
+        self._in = []
+        self._out = []
+        self._d_in = []
+        self._d_out = []
+        self._ascending_order = ascending_order
+        self._freq = defaultdict(int)
+
+    def query(self) -> int:
+        return self._summed if self._ascending_order else -self._summed
+
+    def add(self, x: int) -> None:
+        if not self._ascending_order:
+            x = -x
+
+        self._freq[x] += 1
+        heappush(self._in, -x)
+        self._summed += x
+        self._modify()
+
+    def discard(self, x: int) -> None:
+        if not self._ascending_order:
+            x = -x
+        if self._freq[x] == 0:
+            return
+
+        self._freq[x] -= 1
+
+        if self._in and -self._in[0] == x:
+            self._summed -= x
+            heappop(self._in)
+        elif self._in and -self._in[0] > x:
+            self._summed -= x
+            heappush(self._d_in, -x)
+        else:
+            heappush(self._d_out, x)
+
+        self._modify()
+
+    def set_k(self, k: int) -> None:
+        self._k = k
+        self._modify()
+
+    def get_k(self) -> int:
+        return self._k
+
+    def _modify(self) -> None:
+        while self._out and (len(self._in) - len(self._d_in) < self._k):
+            p = heappop(self._out)
+
+            if self._d_out and p == self._d_out[0]:
+                heappop(self._d_out)
+            else:
+                self._summed += p
+                heappush(self._in, -p)
+
+        while len(self._in) - len(self._d_in) > self._k:
+            p = -heappop(self._in)
+
+            if self._d_in and p == -self._d_in[0]:
+                heappop(self._d_in)
+            else:
+                self._summed -= p
+                heappush(self._out, p)
+
+        while self._d_in and self._in[0] == self._d_in[0]:
+            heappop(self._in)
+            heappop(self._d_in)
+
+    def __len__(self) -> int:
+        return len(self._in) + len(self._out) - len(self._d_in) - len(self._d_out)
+
+    def __contains__(self, x: int) -> bool:
+        if not self._ascending_order:
+            x = -x
+        return self._freq[x] > 0


Test patch:
diff --git a/tests/test_sum_of_top_kth.py b/tests/test_sum_of_top_kth.py
new file mode 100644
index 0000000..aee861e
--- /dev/null
+++ b/tests/test_sum_of_top_kth.py
@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+
+
+from snippets.data_structure.sum_of_top_kth import SumOfTopKth
+
+
+class TestSumOfTopKth:
+    def test_sum_of_top_kth_by_descending_order(self) -> None:
+        s = SumOfTopKth(k=5, ascending_order=False)
+
+        for ai in [3, 1, 4, 1, 5, 9, 2, 6]:
+            s.add(ai)
+
+        actual = s.query()
+        expected = 27
+        assert actual == expected
+
+        s.discard(4)
+        s.add(3)
+        actual = s.query()
+        expected = 26
+        assert actual == expected
+
+        s.set_k(8)
+        actual = s.query()
+        expected = 30
+        assert actual == expected
+
+        s.discard(2)
+        s.add(10)
+        actual = s.query()
+        expected = 38
+        assert actual == expected
+
+    def test_sum_of_top_kth_by_ascending_order(self) -> None:
+        s = SumOfTopKth(k=5, ascending_order=True)
+
+        for ai in [3, 1, 4, 1, 5, 9, 2, 6]:
+            s.add(ai)
+
+        actual = s.query()
+        expected = 11
+        assert actual == expected
+
+        s.discard(4)
+        s.add(5)
+        actual = s.query()
+        expected = 12
+        assert actual == expected
+
+        s.set_k(8)
+        actual = s.query()
+        expected = 32
+        assert actual == expected
+
+        s.discard(6)
+        s.add(2)
+        actual = s.query()
+        expected = 28
+        assert actual == expected