ABC248 D Range Count Query

内容

公式解説と違う解き方なのでメモ。

問題

リストAとクエリ(L,R,X)が与えられる。 Aの区間[L,R]に含まれる 数字X の個数を求める。

atcoder.jp

考えたこと(1-originで書いてます)

• F(L,R) =Aの区間[L,R]に含まれる 数字X の個数とする。
• F(L,R) = F(1,R)-F(1,L-1) である。
  • クエリが、[1,R]と[1,L-1]の二つに分かれる。 　* それぞれを求めて、最後に計算する。
• F(1,q)を求める方法があればいい。
  • これまでに出現した各数字の個数を記憶しておく。(CNT[x]=これまでに出たxの個数)
  • qごとに A[:q]のxを数えるのは計算量的に辛い
    • Aの数字を左から見つつ、クエリ[1,q]のqの小さい順に計算していくことで、Aを一回見るだけで良くなる。

実装

from collections import defaultdict
N = int(input())
A =list(map(int,input().split()))


Q = int(input())
# F[L,R]->F[1,R]-F[1,L-1]

query_sep=[]

for idx in range(Q):
    l,r,x = map(int,input().split())
    query_sep.append((l-1,x,idx))
    query_sep.append((r,x,idx))

#F[1,q]のqの小さいものから見ていく
query_sep = sorted(query_sep,key=lambda x:x[0])


ans=[[] for _ in range(Q)]
CNT=defaultdict(int)

cur_a_idx=0 #Aをどこまで見たか

#F[1,q]のqの小さいものから見ていく
for cur,x,idx in query_sep:

    # A[:q]まで、CNTする
    while  cur_a_idx<cur:
        a = A[cur_a_idx]
        CNT[a]+=1 #A[q]の数を+1する
        cur_a_idx+=1

    #CNTには、A[:q]までのそれぞれの数字の個数が入っている。
    cnt = CNT[x]

    #ansのi-th queryのところに、結果を入れる。
    # L<Rなので、ans[0]はF(1,L) ans[1]はF(1,R)が入る
    ans[idx].append(cnt)


for l in ans:
    #F(L,R)=F(1,R)-F(1,L-1)
    print(l[1]-l[0])