mo

ソースコード

from titan_pylib.algorithm.mo import Mo

view on github

展開済みコード

  1# from titan_pylib.algorithm.mo import Mo
  2from typing import Callable
  3from itertools import chain
  4from math import sqrt, ceil
  5
  6
  7class Mo:
  8    """長さ `n` の列、クエリ数 `q` に対する `Mo's algorithm` です。
  9    :math:`O(\\frac{n}{\\sqrt{q}})` です。
 10
 11    Args:
 12        n (int): 列の長さです。
 13        q (int): クエリの数です。
 14
 15    制約:
 16        :math:`0 \\leq n, 0 \\leq q`
 17    """
 18
 19    def __init__(self, n: int, q: int) -> None:
 20        assert 0 <= n and 0 <= q, f"ValueError: {n=} {q=}"
 21        self.n = n
 22        self.q = q
 23        self.bucket_size = ceil(sqrt(3) * n / sqrt(2 * q)) if q > 0 else n
 24        if self.bucket_size == 0:
 25            self.bucket_size = 1
 26        self.bit = max(n, q).bit_length()
 27        self.msk = (1 << self.bit) - 1
 28        self.bucket = [[] for _ in range(n // self.bucket_size + 1)]
 29        self.cnt = 0
 30
 31    def add_query(self, l: int, r: int) -> None:
 32        """区間 ``[l, r)`` に対するクエリを追加します。
 33        :math:`O(1)` です。
 34
 35        制約:
 36            :math:`0 \\leq l \\leq r \\leq n`
 37        """
 38        assert (
 39            0 <= l <= r <= self.n
 40        ), f"IndexError: {self.__class__.__name__}.add_query({l}, {r}), self.n={self.n}"
 41        self.bucket[l // self.bucket_size].append(
 42            (((r << self.bit) | l) << self.bit) | self.cnt
 43        )
 44        self.cnt += 1
 45
 46    def run(
 47        self,
 48        add: Callable[[int], None],
 49        delete: Callable[[int], None],
 50        out: Callable[[int], None],
 51    ) -> None:
 52        """クエリを実行します。
 53        :math:`O(q\\sqrt{n})` です。
 54
 55        Args:
 56            add (Callable[[int], None]): 引数のインデックスに対応する要素を追加します。
 57            delete (Callable[[int], None]): 引数のインデックスに対応する要素を削除します。
 58            out (Callable[[int], None]): クエリ番号に対する答えを処理します。
 59
 60        制約:
 61            ``q`` 回のクエリを ``add_query`` メソッドで追加する必要があります。
 62        """
 63        assert (
 64            self.cnt == self.q
 65        ), f"Not Enough Queries, now:{self.cnt}, expected:{self.q}"
 66        bucket, bit, msk = self.bucket, self.bit, self.msk
 67        for i, b in enumerate(bucket):
 68            b.sort(reverse=i & 1)
 69        nl, nr = 0, 0
 70        for rli in chain(*bucket):
 71            r, l = rli >> bit >> bit, rli >> bit & msk
 72            while nl > l:
 73                nl -= 1
 74                add(nl)
 75            while nr < r:
 76                add(nr)
 77                nr += 1
 78            while nl < l:
 79                delete(nl)
 80                nl += 1
 81            while nr > r:
 82                nr -= 1
 83                delete(nr)
 84            out(rli & msk)
 85
 86    def runrun(
 87        self,
 88        add_left: Callable[[int], None],
 89        add_right: Callable[[int], None],
 90        delete_left: Callable[[int], None],
 91        delete_right: Callable[[int], None],
 92        out: Callable[[int], None],
 93    ) -> None:
 94        """クエリを実行します。
 95
 96        :math:`O(q\\sqrt{n})` です。
 97
 98        Args:
 99            add_left (Callable[[int], None]): 引数のインデックスに対応する要素を左から追加します。
100            add_right (Callable[[int], None]): 引数のインデックスに対応する要素を右から追加します。
101            delete_left (Callable[[int], None]): 引数のインデックスに対応する要素を左から削除します。
102            delete_right (Callable[[int], None]): 引数のインデックスに対応する要素を右から削除します。
103            out (Callable[[int], None]): クエリ番号に対する答えを処理します。
104
105        制約:
106            ``q`` 回のクエリを ``add_query`` メソッドで追加する必要があります。
107        """
108        assert (
109            self.cnt == self.q
110        ), f"Not Enough Queries, now:{self.cnt}, expected:{self.q}"
111        bucket, bit, msk = self.bucket, self.bit, self.msk
112        for i, b in enumerate(bucket):
113            b.sort(reverse=i & 1)
114        nl, nr = 0, 0
115        for rli in chain(*bucket):
116            r, l = rli >> bit >> bit, rli >> bit & msk
117            while nl > l:
118                nl -= 1
119                add_left(nl)
120            while nr < r:
121                add_right(nr)
122                nr += 1
123            while nl < l:
124                delete_left(nl)
125                nl += 1
126            while nr > r:
127                nr -= 1
128                delete_right(nr)
129            out(rli & msk)

仕様

class Mo(n: int, q: int)[source]

Bases: object

長さ n の列、クエリ数 q に対する Mo’s algorithm です。 \(O(\frac{n}{\sqrt{q}})\) です。

Parameters:

  • n (int) – 列の長さです。

  • q (int) – クエリの数です。

制約:

\(0 \leq n, 0 \leq q\)

add_query(l: int, r: int) → None[source]

区間 [l, r) に対するクエリを追加します。 \(O(1)\) です。

制約:

\(0 \leq l \leq r \leq n\)

run(add: Callable[[int], None], delete: Callable[[int], None], out: Callable[[int], None]) → None[source]

クエリを実行します。 \(O(q\sqrt{n})\) です。

Parameters:

  • add (Callable[[int], None]) – 引数のインデックスに対応する要素を追加します。

  • delete (Callable[[int], None]) – 引数のインデックスに対応する要素を削除します。

  • out (Callable[[int], None]) – クエリ番号に対する答えを処理します。

制約:

q 回のクエリを add_query メソッドで追加する必要があります。

runrun(add_left: Callable[[int], None], add_right: Callable[[int], None], delete_left: Callable[[int], None], delete_right: Callable[[int], None], out: Callable[[int], None]) → None[source]

クエリを実行します。

\(O(q\sqrt{n})\) です。

Parameters:

  • add_left (Callable[[int], None]) – 引数のインデックスに対応する要素を左から追加します。

  • add_right (Callable[[int], None]) – 引数のインデックスに対応する要素を右から追加します。

  • delete_left (Callable[[int], None]) – 引数のインデックスに対応する要素を左から削除します。

  • delete_right (Callable[[int], None]) – 引数のインデックスに対応する要素を右から削除します。

  • out (Callable[[int], None]) – クエリ番号に対する答えを処理します。

制約:

q 回のクエリを add_query メソッドで追加する必要があります。