一定区間の整数の出現回数を数え上げる方法 4種類 （imos法を含む）[Python実装]

背景

競技プログラミングAtCoderのABC127のC問題を解いていた際に、連続した整数の集合の共通部分がいくつあるかを考える問題を考えました。様々な解法があった問題で非常に面白いと感じましたので、記事にすることにしました。問題は以下のリンクから参照することができます。ABC127 C問題の解説も行っていますので、ネタバレ注意です。

AtCoder Beginner Contest 127 C - Prison

C - Prison

まずは問題文から見ていくことにしましょう。

問題文の概略を書きます。
M個のゲートがあり、そのゲートを開けるためにIDカードが必要です。N枚のIDカードがあり１からNまでの番号がふられています。あるゲートはLi以上Ri以下の番号のIDカードのみ開けることが可能です。全てのゲートを開けることができるIDカードは幾つ存在しているかを数える問題です。

まず入力例で考えてみますと、
４枚のIDカードがあり、それぞれ１、２、３、４の番号がふられています。
３個のゲートがあり、１番目のゲートは１、２、３のIDカードで２番目のゲートは２、３、４のIDカードで開けることができます。
全てのゲートを開けることができるのは２と３のIDカードであることがわかります。つまり、全てのゲートを開けることのできるIDカードの積集合を取りたいという考えになりました。

setを使って集合を考えた　方針１

結論から言いますと、この方法では計算時間が膨大でTLEになって通過しませんでした。

方針としては１つののゲートを開けることができるIDカードの番号の最小値と最大値の情報が1行で入力されるので、最小値から最大値までの数をrangeで作成し、set型にに変換しました。その操作を全てのゲートに対して行い、全ての集合で積集合を取りました。

TLEになった原因を考えます。１個のゲートを開けるIDカードの個数が最大で${10^{5}}$個もあるので、毎回${10^{5}}$個のデータをsetに変換し、M個（最大で${10^{5}}$）のゲートに対して行うことは非常に計算量が多いことがわかります。その作業をM回行うので、最悪の場合${O(NM)}$くらいかかると見積もりました。制約から計算量が${O(10^{10})}$でTLEになりました。

参考程度にpythonのコードを載せます。

n, m = map(int, input().split())
line = set(range(1, n + 1))

for i in range(m):
    l, r = map(int, input().split())
    line &= set(range(l, r + 1))

print(len(line))

方針１を修正　方針２

集合を考える場合はIDカードが連続していないような状況では有効になるかも知れないですが、今回の場合は計算量も多く有効ではありませんでした。計算量が最大でも${O(N)}$程度にしなければならないことがわかります。

そこでもう一度方針を考えるとデータが連続しているので、IDカードの最小値Liのうち最大のIDカードの番号と最大値Riのうち最小の番号を保存すればその間の数が共通するIDカードの個数ではないかと考えました。

この方針が思いついた背景は、ゲートを開けるIDカードの数字をゲート毎に並べて羅列していくことで気付きました。
今回の場合はIDカードが連続して入力されるので、この方法で良いです。Ri - Li + 1個のIDカードで開けることができます。

ただ、注意する点としてIDカードの範囲が重ならない場合にRi - Li + 1の値が負になってしまうので、負の場合は条件を満たすIDカードは存在しないので０を出力する実装を行わなければなりません。

n, m = map(int, input().split())
l_max = 1
r_min = 10 ** 5

for i in range(m):
    l, r = map(int, input().split())
    l_max = max(l, l_max)
    r_min = min(r, r_min)

print(max(0, r_min - l_max + 1))

このコードでACを取ることができました。ACを取るだけならこれで満足ですが、解説放送を見ていた際に、コメントでimos法で解いたと言うコメントがありまして、興味が出て別解を再考しました。

imos法に入る前に　方針３

まずはimos法に取り組む前に、imos法と似たアプローチである方法で考えてみたいと思います。
結論を先に述べますと、この方法では計算量が多くTLEになってしまうと思います。それを解決するのが我らがimos法です。

それでは思考を整理していきます。
IDカードが開けることができるゲートの数を考えた際に、全てのゲートを開けるIDカードはどのゲートも開けることができるので、M個になります。M未満の場合は開けることができないゲートが存在します。
つまり、M個のゲートを開けることができるIDカードの数を数えれば良いことがわかります。

それでは実装してみたいと思います。まず、IDカードの番号が何個のゲートを開けることができるか保存するための配列を用意します。cardid_2_num_mapという配列を用意しました。要素数はIDカードの最大値から1つ多くとっています。この意図は入力される数字が１以上N以下なのでindexを指定する際にそのまま使用することでエラーを減らそうとしました。０番目の配列は使うことはありません。

カードIDの番号がindexで0から${10^5}$まで、配列の中身がそのIDカードがゲートを開けられる数を保持した配列を考える際に、[0, 1, 2, 2, 1, 0, ...]となることが入力例の場合に得たい配列です。

forが二重になっていることを考えても、計算量が多そうです。計算量のオーダーは最大で${O(NM)}$となり、積集合を考えた際と同様にTLEしてしまいます。この問題を解決したのがimos法です。次のセクションで説明していきます。

n, m = map(int, input().split())
cardid_2_num_map = [0] * (10**5 + 1)

for i in range(m):
    l, r = map(int, input().split())
    for j in range(l, r + 1):
        cardid_2_num_map[j] += 1

max_sum_num = max(cardid_2_num_map)
count_max_card = cardid_2_num_map.count(max_sum_num)

if max_sum_num == m:
    print(count_max_card)
else:
    print(0)

imos法　今回の主役　方針４

参考にさせていただいた記事を載せます。

いもす法 - いもす研 (imos laboratory)
imos法は連続した数の出現回数を数える際に累積和を用いて、毎回数え上げていた処理を１回分にまとめることで、計算量を減らすことに成功したアルゴリズムです。

それでは上の問題に戻って考えていきます。IDカードがゲートを開けることができる数を数え上げて、配列内の値を更新していきました。計算量のオーダーは最大で${O(NM)}$となりTLEになりました。オーダーを${O(M)}$程度に抑えることができれば解決できそうです。

そこで、入力例に戻って考えますと４枚のIDカードがあり、それぞれ１、２、３、４の番号がふられています。
３個のゲートがあり、１番目のゲートは１、２、３のIDカードで２番目のゲートは２、３、４のIDカードで開けることができます。

方針３と同様に、カードIDの番号がindexで0, 1, 2, 3, 4, ..., ${10^5}$で、配列の中身がそのIDカードがゲートを開けられる数を保持した配列を考える際に、[0, 1, 2, 2, 1, 0, ...]となることが入力例の場合に得たい配列です。０番目のindexはカードIDが０を示すので使うことはありません。このような配列を得る際に、imos法を使います。

ます、配列を要素数${10^5 + 2}$個、値を０で初期化します。＋２の意図は後ほど説明します。
１番目のゲートは１、２、３のカードで開けることができるのでindex番号１に＋１、index番号４に−１を行います。
[0, 1, 0, 0, -1, 0, ..., 0]
同様に２番目のゲートは２、３、４のIDカードで開けることができるので、index番号２に＋１、index番号５に−１を行います。
[0, 1, 1, 0, -1, -1, ..., 0]
IDカードの右端＋１のindexに対して−１の操作するので、配列の要素数を${10^5}$にさらに１加えて＋２しました（＋１はindex番号と入力を合わせるため行いました）。

次に、この配列をindex番号０から最後まで累積和をとっていきます。
A[i + 1] = A[i] + A[i + 1]
を行います。結果として方針３と同様に目的とした配列[0, 1, 2, 2, 1, 0, ...]を得ることができます。

そして、配列の最大値がゲートの個数に等しければ全てのゲートを開けることができるIDカードが存在して、その最大値のIDカードの枚数を数えます。計算量のオーダーは最大で${O(M)}$程度となりACすることができました。

n, m = map(int, input().split())
cardid_2_num_map = [0] * (10**5 + 2)

for i in range(m):
    l, r = map(int, input().split())
    cardid_2_num_map[l] += 1
    cardid_2_num_map[r + 1] -= 1

for i in range(1, n + 1):
    cardid_2_num_map[i] += cardid_2_num_map[i - 1]

max_sum_num = max(cardid_2_num_map)
count_max_card = cardid_2_num_map.count(max_sum_num)

if max_sum_num == m:
    print(count_max_card)
else:
    print(0)

考察

数えたい範囲の左端のindexに＋１、右端＋１のindexに−１を行うことで最終的に累積和を取る段階で左端から右端までだけに＋１を行うことが可能になり、全ての領域に対して１回で出現回数を数えることができます。

x < 左端　の領域は何もしない
左端 <= x <= 右端　の領域は＋１されて
右端 < x　の領域は何もしないので

意図した数え上げを行うことができています。

この思考法は多次元にも応用でき計算量を削減できるみたいなので、非常に興味深いと感じました。参考にさせていただいた記事のイラストが非常に理解しやすかったです。以上にします。