トップ > AtCoder > ABC234D - Prefix K-th Max

ABC234D - Prefix K-th Max

AtCoder AtCoder-Beginner Contest Binary Indexed Tree (BIT) 優先度付きキュー Python

優先度付きキュー

考え方

$K$要素のリストを管理し，
  • 「(リストの中で一番小さい値) < (新しい要素の値)」なら，値を置き換え，
  • そうでないなら，リストを変更しない
とすれば，このリストの最小値が各ステップで求めたい値となる．

これは，最初の要素が常に最小要素となる，優先度付きキュー（heapq）で実現できる．
【参考】heapq --- ヒープキューアルゴリズム — Python 3.10.0b2 ドキュメント

回答例

import heapq

N, K = map(int, input().split())
P = list(map(int, input().split()))

q = P[:K]
heapq.heapify(q)
print(q[0])

for p in P[K:]:
  if q[0] < p:
    heapq.heappop(q)
    heapq.heappush(q, p)
  print(q[0])

参考

BIT

【参考】Binary Indexed Tree (BIT) / Fenwick Tree - 競プロはじめました

考え方

ソートされた状態でリストが管理し，大きい方から$K$番目（小さい方から$i - K + 1$番目）の値を求めたい．

これは，

  • BITで（空想の）リストに$i$が入っているか入っていないかを 0 / 1 で管理し，
  • 各ステップで，小さい方から和をとったときに$i - K + 1$となるindex（＝空想のリストの要素の値）を二分探索で求める
ことで実現できる．

回答例

フツーの二分探索をすると以下．BIT上の二分探索をすると高速化できる（後述）．

N, K = map(int, input().split())
P = list(map(int, input().split()))

bit_size = N + 1
bit = [0] * bit_size

def add(i, x):
  while i < bit_size:
      bit[i] += x
      i += i & -i

def get(i):
  ans = 0
  while i > 0:
      ans += bit[i]
      i -= i & -i
  return ans
    
for i in range(K - 1):
  add(P[i], 1)
  
for i in range(K - 1, N):
  add(P[i], 1)
  l, r = 0, N
  while r - l > 1:
    mid = (r + l) // 2
    if get(mid) < i + 1 - K + 1:
      l = mid
    else:
      r = mid
  print(r)


BIT上の二分探索をすることで，高速化する．

N, K = map(int, input().split())
P = list(map(int, input().split()))

bit_size = N + 1
bit = [0] * bit_size

def add(i, x):
  while i < bit_size:
      bit[i] += x
      i += i & -i

def get(i):
  ans = 0
  while i > 0:
      ans += bit[i]
      i -= i & -i
  return ans

# https://ikatakos.com/pot/programming_algorithm/data_structure/binary_indexed_tree
def lower_bound(x):
  """ 累積和がx以上になる最小のindexと、その直前までの累積和 """
  depth = N.bit_length()
  pos, acc = 0, 0
  for i in range(depth, -1, -1):
    k = pos + (1 << i)
    if k <= N and acc + bit[k] < x:
      acc += bit[k]
      pos += 1 << i
  return pos + 1, acc
    
for i in range(K - 1):
  add(P[i], 1)
  
for i in range(K - 1, N):
  add(P[i], 1)
  ans, acc = lower_bound(i + 1 - K + 1)
  print(ans)