トップ > AtCoder > ABC231F - Jealous Two

ABC231F - Jealous Two

AtCoder AtCoder-Beginner Contest Python Binary Indexed Tree (BIT)


【関連】

考え方

$A$の昇順・$B$の降順で見ていく．絵を書くとわかりやすい．

各$B_{i}$に対し，

  • bit[i] += 1とし，
  • ans +=$\sum_{j (\geq i)} \mathrm{bit\,} [j]$としていく
ことで答えが求まる．

これは，BIT (Binary Indexed Tree (BIT) / Fenwick Tree - 競プロはじめました) を使えばできる．

今回実装したBITは$\sum_{j (\geq i)} \mathrm{bit\,} [j]$ではなく，$\sum_{j (\textcolor{red}{\leq} i)} \mathrm{bit\,} [j]$を計算するものなので，$B_{i}$の個数は$\mathrm{bit\,} [N - 1 - i]$で管理した．


また，$(A_{i}, B_{i}) = (A_{j}, B_{j})$があると，その分を余計にカウントしないといけないことに注意する．

ABC231F

回答例

N = int(input())
A = list(map(int, input().split()))
B = list(map(int, input().split()))

Xdict = {x:i for i, x in enumerate(sorted(list(set(A))))}
Ydict = {x:i for i, x in enumerate(sorted(list(set(B))))}
pos = [(Xdict[a], Ydict[b]) for a, b in zip(A, B)]
pos.sort(lambda x: (x[0], -x[1]))

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

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

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

ans = 0
cnt = 0
lasta, lastb = -1, -1
for a, b in pos:
  b = N - 1 - b
  add(b, 1)
  ans += get(b)
  if lasta == a and lastb == b:
    cnt += 1
    ans += cnt
  else:
    lasta, lastb = a, b
    cnt = 0

print(ans)

【参考】【競プロ実況】ABC231 F問題【かつっぱ】 - YouTube

参考