ABC 140 F - Many Slimes【Python】

問題はこちら

大きいスライムから順に貪欲につくるのが最適のはずと決め打って愚直にシミュレーションした。高速に処理するには平衡二分木と優先度付きキューがあればよい。C++ならmultisetとpriority_queueでいいのだろうが、pythonにはheapqはあるけど平衡二分木がないのでやや性能は劣るがBITで代用。解説も読んだけど意味が分からなかった。

import sys
from heapq import heappop, heappush
readline = sys.stdin.readline


class BIT():
 def __init__(self, n):
   self.n = n
   self.tree = [0] * (n + 1)

 def add(self, i, x):
   while i <= self.n:
     self.tree[i] += x
     i += i & -i

 def sum_(self, i):
   ret = 0
   while i > 0:
     ret += self.tree[i]
     i -= i & -i
   return ret

 def pop(self, i):
   if i == 0:
     return 0
   sum_ = self.sum_(i - 1)
   ng, ok = -1, i - 1
   while ok - ng > 1:
     mid = (ok + ng) // 2
     if self.sum_(mid) >= sum_:
       ok = mid
     else:
       ng = mid
   if ok > 0:
     self.add(ok, -1)
   return ok


def main():
 N = int(readline())
 S = list(map(int, readline().split()))

 S_set = list(set(S)); S_set.sort()
 M = len(S_set)
 S_to_I = {s: i + 1 for i, s in enumerate(S_set)}
 S_mul = BIT(M)
 for s in S:
   i = S_to_I[s]
   S_mul.add(i, 1)

 G = []
 heappush(G, -S_mul.pop(M + 1))
 for i in range(N):
   tmp = G.copy()
   for _ in range(1 << i):
     x = -heappop(tmp)
     j = S_mul.pop(x)
     if j == 0:
       print('No')
       return
     heappush(G, -j)
 print('Yes')


if __name__ == '__main__':
 main()

この記事が気に入ったらサポートをしてみませんか?