Go to Qiita Advent Calendar 2024 Top
LoginSignup
1
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@alumite14

特定要素の区間内個数を高速に求める方法（更新OK）

Last updated at Posted at 2024-09-10

はじめましての人ははじめまして。alumiと申します。
区間内個数のライブラリで遊んでいたら現行の手法の改良を思いつき、実際に定数倍高速化できたので紹介します。

静的配列に対するクエリ

以下制約は次の値と仮定します。

  • $0 \le N \le 5 \times 10^5$
  • $0 \le Q \le 5 \times 10^5$
  • $0 \le a_i \le 10^9$
  • $0 \le l \le r \le N$
  • $0 \le x \le 10^9$

改良前の実装（dict + 二分探索）

静的配列$A$が最初に与えられ、その後に$Q$個のクエリ$l,r,x$が与えられる状況を考えます。
各クエリにおいて$A_l,A_{l+1},...,A_{r-1}$のうち$x$と等しいもの数を答えることがタスクです。

naiveな方法では$A$の$[l,r)$の要素を一つずつ$x$と比較する必要があるので、時間計算量は$O(N)$となります。
そこで任意要素に対して、その要素が出現するindexを昇順に列挙した辞書$index\_map$を持つようにします。具体的には、

  • A = [3, 3, 0, 2, 2, 2, 3, 1, 3, 1, 3, 0, 5, 4, 1]に対して、
  • index_map = {3: [0, 1, 6, 8, 10], 0: [2, 11], 2: [3, 4, 5], 1: [7, 9, 14], 5: [12], 4: [13]}
    のような辞書を構成するようにします。

そうすると各クエリ$l,r,x$に対して、

  • 辞書に含まれていない$x$なら0を返す（$O(1)$）
  • 辞書に含まれている場合
    • index_map[x]で$x$が存在する位置だけを取得
    • bisect_left(index_map[x], r)でindexが$r-1$以下となる要素数を取得
    • bisect_left(index_map[x], l)でindexが$l-1$以下となる要素数を取得
    • これらの差が$A[l:r]$にある$x$の答えとなる（$O(\log{N})$）
      というふうにタスクをこなせます。

具体的に見てみましょう。前述のようにA = [3, 3, 0, 2, 2, 2, 3, 1, 3, 1, 3, 0, 5, 4, 1]のとき、index_map = {3: [0, 1, 6, 8, 10], 0: [2, 11], 2: [3, 4, 5], 1: [7, 9, 14], 5: [12], 4: [13]}となります。このとき、

  • $l,r,x=2,8,4$
    • 辞書に$4$はないので$0$を返す
  • $l,r,x=4,9,3$
    • index_map[3] = [0, 1, 6, 8, 10]で
    • bisect_left(index_map[3], 9) = 4（[0, 1, 6, 8, (10)]）
    • bisect_left(index_map[3], 4) = 2（[0, 1, (6), 8, 10]）
    • よって$4-2=2$を返す
      • 実際A[4:9] = [2, 2, 3, 1, 3]で$3$は2個
  • $l,r,x=0,3,2$
    • index_map[2] = [3, 4, 5]で
    • bisect_left(index_map[2], 3) = 0（[(3), 4, 5]）
    • bisect_left(index_map[2], 0) = 0（[(3), 4, 5]）
    • よって$0-0=0$を返す
      • 実際A[0:3] = [3, 3, 0]で$2$は0個

上の考えをクラスにしたものが以下の実装です。一回のクエリについて2回の二分探索を行っているため、1クエリあたりの計算量は$O(\log{N})$となります。

from collections import defaultdict
from bisect import bisect_left

class StaticRangeCounter:
    def __init__(self, init_values: list) -> None:
        self.index_map = defaultdict(list)
        for i, x in enumerate(init_values):
            self.index_map[x].append(i)

    def count_range(self, l: int, r: int, x: int) -> int:
        """Count the frequency of x in [l, r)."""
        li = bisect_left(self.index_map[x], l)
        ri = bisect_left(self.index_map[x], r)
        return ri - li

改良実装（list + 二分探索）

上記の実装でも問題ない場合のほうが多いのですが、速度を考えると辞書の参照はlistと比べてオーバーヘッドが大きいので避けたいところです。そこで1次元listを利用したいので、「任意要素に対するindexが、対応する位置に昇順に並べられないかどうか」を考えてみます。つまり

  • A = [3, 3, 0, 2, 2, 2, 3, 1, 3, 1, 3, 0, 5, 4, 1]に対して、
  • index_map = {3: [0, 1, 6, 8, 10], 0: [2, 11], 2: [3, 4, 5], 1: [7, 9, 14], 5: [12], 4: [13]}だったのを
  • indices = [|2, 11|, |7, 9, 14|, |3, 4, 5|, |0, 1, 6, 8, 10|, |13|, |12|]

とすることを考えます。当然これは昇順でもなければ、各要素の最初・最後のindexの位置も特定できないので二分探索できません。どうすればいいでしょうか。

要は$A$の要素$x,y$が$x<y$のとき「$x$の位置のindexを表す値」$<$「$y$の位置のindexを表す値」となっていて、かつ適当な変換によって$l,r$を$\min($$x$の位置のindexを表す値$) \le l,r \le \max($$x$の位置のindexを表す値$)$とできれば、二分探索できそうです。
上記を実現する変換は実にシンプルです。つまり$A$の要素数を$N$とすると「$A[i] \times N + i$」で変換できます。
実際の値がどのようにmappingされるかを見ると理解しやすいと思います。つまり、

  • $A[i]=0$→変換後は$[0,N)$
  • $A[i]=1$→変換後は$[N,2N)$
  • ...
  • $A[i]=x$→変換後は$[xN,(x+1)N)$
  • ...
    と互いにダブらない値にmappingされます。実際の値で見てみると
  • N = 15, A = [3, 3, 0, 2, 2, 2, 3, 1, 3, 1, 3, 0, 5, 4, 1]なので
    • $A[i]=0$ → 変換後は$[0,15)$
    • $A[i]=1$ → 変換後は$[15,30)$
    • $A[i]=2$ → 変換後は$[30,45)$
    • $A[i]=3$ → 変換後は$[45,60)$
    • $A[i]=4$ → 変換後は$[60,75)$
    • $A[i]=5$ → 変換後は$[75,90)$
  • とmappingされ、実際に変換すると
  • cA = [45, 46, 2, 33, 34, 35, 51, 22, 53, 24, 55, 11, 87, 73, 29]となります。
    そしてcAをソートすると、
  • cA = [|2, 11|, |22, 24, 29|, |33, 34, 35|, |45, 46, 51, 53, 55|, |73|, |87|]
    となります。仕切りは変換前同じ数だったグループとなるよう入れています。前述のindicesの割り振りと各値の値域の対応が一致しているのがわかるでしょうか。

そしてクエリとして与えられた$l,r,x$を$x \times N + l$, $x \times N + r$と変換してやれば、これらの値は必ず$x$の変換後の値域$[xN, (x+1)N)$に収まるので、最初の実装と同様に二分探索できます。わーい。

ここまでを実装すると次のようになります。

class StaticRangeCounter:
    def __init__(self, arr) -> None:
        self.size = len(arr)
        self.indices = sorted(a * self.size + i for i, a in enumerate(arr))

    def range_count(self, l: int, r: int, x: int) -> int:
        """Count the frequency of x in [l, r)."""
        li = bisect_left(self.indices, x * self.size + l)
        ri = bisect_left(self.indices, x * self.size + r)
        return ri - li

これと同様の発想のコードでLibrary CheckerのStatic Range Frequencyでpypy最速を出せました（2024-09-10現在、Submission #228093）。

動的配列に対するクエリ

改良前の実装（dict + SortedSet + 二分探索）

動的配列では静的配列でのタスクに加えて、$A[i]$の要素を指定要素$x$に置き換える更新クエリ$i, x$が追加されます。
静的配列における改良前の実装で、適切に$index\_map$の値を更新できれば、他の実装は変えなくても良さそうです。
$index\_map$では特定要素のindexをlistで持っていましたね。よって更新によってある要素を消した場合、list内の要素の削除を実行することになるので$O(N)$かかります。追加に関しても同様です。

まずこのコストを無視して、更新に必要な情報・操作を整理しましょう。更新には

  • 何を消すか=現在の$A[i]$の値を取得
    • 消す要素の$index\_map$から$i$を削除（$O(N)$）
    • 追加要素$x$の$index\_map$に$i$を昇順を保ったまま追加（$O(N)$）
  • $A[i] = x$とする
    という操作が必要です。太字の条件はcountクエリで二分探索するために必要でしたね。

よって$index\_map$をlistから、「任意要素の追加・削除を高速にできて、要素を昇順に持ってくれる」データ構造で置き換えれば、更新クエリに対応できそうです。

このような操作に都合のいいデータ構造としてSortedSetというものがあります。これは一意の値を昇順に格納できるデータ型であり、

  • 要素の追加・削除
  • 任意要素の位置の特定

が高速1で可能となっています。便利～

静的配列の改良前の実装で、listで初期化していた部分をSortedSetで置き換えることで、動的配列に対しても高速に対応できる実装が得られます。わーい

Library Checkerではpypyのversionが3.9のためsortedcontainersライブラリが使えません。tatyam様のSortedSetなどで代替してください：Submission #234471

from sortedcontainers import SortedList
from collections import defaultdict

class DynamicRangeCounter:
    def __init__(self, init_values: list) -> None:
        self.values = init_values
+       self.index_map = defaultdict(SortedList)
        for i, v in enumerate(init_values):
            self.index_map[v].add(i)

+   def update(self, i: int, x) -> None:
+       """Update i to x (0-indexed)"""
+       old_val = self.values[i]
+       new_val = x
+
+       self.index_map[old_val].remove(i)
+       if not self.index_map[old_val]:  # If the list becomes empty
+           del self.index_map[old_val]  # Remove the key from the dictionary
+
+       self.index_map[new_val].add(i)
+       self.values[i] = new_val

    def range_count(self, l: int, r: int, x) -> int:
        """Count the frequency of x in [l, r)."""
        if l >= r or x not in self.index_map: return 0
+       li = self.index_map[x].bisect_left(l)
+       ri = self.index_map[x].bisect_left(r)
        return ri - li

改良実装（SortedSet + 二分探索）

勘の良い方ならお気づきかもしれませんが、上記の動的配列の実装は、静的配列の改良と全く同様に高速化ができます。updateメソッドは空になった辞書を消す工程がなくなっているので少し簡素になっています。

class DynamicRangeCounter:
    def __init__(self, init_values: list) -> None:
        self.n = len(init_values)
        self.values = init_values
        self.indices = SortedSet()
        for i, v in enumerate(init_values):
            self.indices.add(v * self.n + i)

    def update(self, i: int, x) -> None:
        """Update i to x (0-indexed)"""
        old_val = self.values[i]
        new_val = x
        self.indices.discard(old_val * self.n + i)
        self.indices.add(new_val * self.n + i)
        self.values[i] = new_val

    def range_count(self, l: int, r: int, x) -> int:
        """Count the frequency of x in [l, r)."""
        if l >= r: return 0
        li = self.indices.index(x * self.n + l)
        ri = self.indices.index(x * self.n + r)
        return ri - li

これと同様の発想のコードでLibrary CheckerのPoint Set Range Frequencyでpypy最速を出せました（2024-09-10現在、Submission #234476）。おそらくSortedSetの実装次第でだいぶ変わると思うので、より良い実装を思い付かれたら私の屍を超えていってください。

結び

というわけで区間内個数を高速に求める方法に関する解説と実装でした。まだコンテストで使ったことがないので改良点や追加機能など特に思いつきませんが、こうしたほうが良いよ！ってのがあれば教えてくれるととても嬉しいです。
よかったらいいね押してね。またね。

Verification

References

  1. ここの時間計算量はSortedSetの実装に依ります。平方分割で実装すれば$O(\sqrt{N})$ですし、スキップリストや二分木で実装すれば$O(\log{N})$になります。 ↩

1
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
1
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?