ABC330 E Mex and Update をPythonで解く【Atcoder】

問題はこちらです。 atcoder.jp

Pythonでセグ木を使う解法記事がないので書いておきます。 まず、セグ木でAに0~N+1がそれぞれ登場する回数を管理します。
クエリが飛んできた時には、のセグ木上の登場回数1へらし、xを1増やします。（実装ではupdateメソッドに該当します）
次にMEXを求める処理を考えます。
ここで、MEXは登場回数が0回である最小の数と言い換えることができます。
したがって、セグ木を二分探索すればよいです。

# oj t -c "python3 main.py"
import sys,math
from collections import defaultdict,deque
from itertools import combinations,permutations,accumulate,product
from bisect import bisect,bisect_left,bisect_right
from heapq import heappop,heappush,heapify
from sortedcontainers import SortedList,SortedSet
def input(): return sys.stdin.readline().rstrip()
def ii(): return int(input())
def ms(): return map(int, input().split())
def li(): return list(map(int,input().split()))
inf = pow(10,18)
#////////////////////////////////////
class segtree:  # 0-index
  # 要素treeの初期化（最初はinit_val）treeがセグメント木そのもの
  def __init__(self,init_val,segfunc,ide_ele):
    self.segfunc = segfunc
    self.ide_ele = ide_ele
    self.n = len(init_val)
    self.size = 1
    while self.size < self.n:
      self.size *= 2
    self.tree = [self.ide_ele] * (self.size * 2)
    for i in range(self.n):
      self.tree[self.size + i] = init_val[i]
    for i in range(self.size-1,0,-1):
      self.tree[i] = self.segfunc(self.tree[2*i], self.tree[2*i+1])

  def get(self,pos): # 元の配列のpos番目の要素を取得する(0-index)
    return self.tree[pos+self.size]

  def update(self, pos, x): 
    if pos<=N+1:
      pos += self.size
      self.tree[pos] -= 1
      while pos >= 2:
        pos //= 2
        self.tree[pos] = self.segfunc(self.tree[pos * 2], self.tree[pos * 2 + 1])
    if x<=N+1:
      x += self.size
      self.tree[x] += 1
      while x >= 2:
        x //= 2
        self.tree[x] = self.segfunc(self.tree[x * 2], self.tree[x * 2 + 1])

  # [l, r) の最大値を求める処理
  def query(self, l, r):
    ret = self.ide_ele
    l += self.size
    r += self.size

    while l < r:
      if l & 1:
        ret = self.segfunc(ret, self.tree[l])
        l += 1
      if r & 1:
        ret = self.segfunc(ret, self.tree[r-1])
      l >>= 1
      r >>= 1
    return ret
# 最小値
def segfunc(x, y): return min(x,y)
ide_ele = pow(10,18)

N,Q = ms()
A = li()
B = [0]*(N+1)
for a in A: 
  if a<=N: B[a] += 1
  
seg = segtree(B,segfunc,ide_ele)
for _ in range(Q):
  i,x = ms()
  tmp = A[i-1]
  seg.update(tmp,x)
  A[i-1] = x
  ng,ok = 0,N+2
  while abs(ok-ng)>1:
    mid = (ok+ng)//2
    if seg.query(0,mid)==0: # 条件を満たすなら
      ok = mid
    else:
      ng = mid
  print(ok-1)

AHC030をPythonでゆっくりやってみる

AHC030のベイズ推定を使う解法をPythonでやってみます。 以下の記事を参考にさせていただきました。 qiita.com

M=2のときの配置推定

まずはM=2のときの配置推定（ベイズの定理）からやってみます。

ベイズの定理の細かい内容は飛ばします。（僕もお気持ち程度しかわかっていません）

とりあえず、油田の配置可能集合とし、これまでマス集合に対して占いをし結果がであった時のが知りたいです。

つまり、真の配置がであるとき、占い結果がになる確率を計算します。

これを計算するには、「マスの集合を占い、真の石油埋蔵量がのときに占い結果がとなる確率」を考えればよいです。

したがって、真の配置がのときのマス集合の石油埋蔵量をと書くと、となります。

実装時には、占い結果が近似値で入力されることに注意が必要です。

例えば、占い結果が1.3のときジャッジプログラムからは1が返されます。

そのため、確率を考えるときは～の範囲を考える必要があります。

import math
from random import randint

N,M,E = map(float, input().split())
N,M = int(N),int(M)
D = list()
SUM_D = 0
IJ = list()
MIN_D = 10**18
for _ in range(M):
  d,*ij = map(int,input().split())
  MIN_D = min(MIN_D,d)
  D.append(d)
  SUM_D += d
  L = list()
  for ind in range(0,2*d,2):
    L.append((ij[ind],ij[ind+1]))
  IJ.append(L)
for ind in range(M): # 各油田が左上から順に並ぶようにする
  IJ[ind].sort()

def make_all_board(): # 配置可能な油田位置を列挙する（ここはもっといい書き方があると思います）
  BOARD = list()
  m_board = list() # [1番目の油田の配置可能位置,2番目の油田の配置可能位置]
  for m in range(M):
    tmp = list()
    for i in range(N):
      for j in range(N): # m番目の油田の左上を(i,j)に合わせる
        nn = set()
        ng = False
        for ii,jj in IJ[m]:
          ni,nj = i+ii,j+jj
          if not (0<=ni<N and 0<=nj<N): 
            ng = True
            break
          nn.add((ni,nj))
        if ng: continue
        tmp.append(nn)
    m_board.append(tmp)
  for m1 in range(len(m_board[0])):
    for m2 in range(len(m_board[1])):
      nn = [[0]*N for _ in range(N)]
      for i,j in m_board[0][m1]:
        nn[i][j] += 1
      for i,j in m_board[1][m2]:
        nn[i][j] += 1
      BOARD.append(nn)
  
  return BOARD

def normal_cdf(x,mean,sigma):
  return 0.5 * (1.0 + math.erf((x - mean) / (sigma * math.sqrt(2))))

def probability_in_range(l,r,mean,sigma):
  assert l <= r
  if mean < l:
    return probability_in_range(2.0 * mean - r, 2.0 * mean - l, mean, sigma)
  p_l = normal_cdf(l, mean, sigma)
  p_r = normal_cdf(r, mean, sigma)
  return p_r - p_l

def q(k,s,r): # kマスの集合Sを占い、石油埋蔵量がsのとき、占い結果がrとなる確率を返す
  # 問題文に従って平均,分散を計算
  mean = (k-s)*E + s*(1-E)
  sigma = (k*E*(1-E))**0.5
  
  # 占い結果は近似値なので、+-0.5の範囲を調べる
  if r==0: return probability_in_range(-1e10, r+0.5, mean, sigma) # 0の時だけ-∞にする必要がある
  else: return probability_in_range(r-0.5, r+0.5, mean, sigma)
  
def normalize(p): # 事後確率の計算
  SUM = sum(p)
  assert SUM>0
  return [x/SUM for x in p]

def bayesian_inference(BOARD):
  n = len(BOARD)
  P = [1.0/n]*n
  while True:
    # 占いする座標を20個ランダムに選ぶ
    S = set()
    while len(S)<20:
      i,j = randint(0,N-1),randint(0,N-1)
      S.add((i,j))
    
    # 占いをする
    print('q',len(S),end=' ')
    for i,j in S: print(i,j,end=' ')
    print(flush=True)
    r = int(input())
    
    for ind in range(n): 
      # i番目の盤面の時のv(s)
      vs = 0
      for i,j in S:
        vs += BOARD[ind][i][j]
      
      # 尤度をかける
      P[ind] *= q(len(S),vs,r)
      
    P = normalize(P)
    
    max_p = max(P)
    max_ind = P.index(max_p)
    
    # 確率が80%以上なら解答する
    if max_p>0.8:
      ANS = set()
      for i in range(N):
        for j in range(N):
          if BOARD[max_ind][i][j]>=1:
            ANS.add((i,j))
      
      print('a',len(ANS),end=' ')
      for i,j in ANS: print(i,j,end=' ')
      print(flush=True)
      r = int(input())
      if r==1: exit() # 正解なら終了
      else: P[max_ind] = 0.0 # 不正解なら事前確率を0にして終了
      
bayesian_inference(make_all_board())

このコードをseed=0で試してみると、cost=2.459675になりました。 占いをランダムにしているのにここまで改善するとは驚きです。

占いマス生成（相互情報量）

気が向いたら続きをやりたいと思います、、、

ABC268 D Unique Username をPythonで解く【AtCoder】

問題はこちらです atcoder.jp

全列挙が間に合いそうなので、考えられる文字列を全列挙します。
全列挙にはDFSを使います（典型ですね）。
Sの並び替えはpermutationsを使い、"_"の追加をDFSでやります。
実装して以下のコードを提出してみます。

import sys,math
import pypyjit
pypyjit.set_param('max_unroll_recursion=-1')
sys.setrecursionlimit(10**7)
from collections import defaultdict,deque
from itertools import combinations,permutations,accumulate,product
from bisect import bisect,bisect_left,bisect_right
from heapq import heappop,heappush,heapify
#from sortedcontainers import SortedList,SortedSet
def input(): return sys.stdin.readline().rstrip()
def ii(): return int(input())
def ms(): return map(int, input().split())
def li(): return list(map(int,input().split()))
#////////////////////////////
N,M = ms()
S = [input() for _ in range(N)]
T = set()
for _ in range(M):
  t = input()
  T.add(t)
def dfs(s,ind):
  s += per[ind]
  if len(s)>16: return
  if ind==N-1:
    if not s in T and 16>=len(s)>=3:
      print(s)
      exit()
    return
  for i in range(1,16-len(s)+1):
    tmp = s + "_"*i
    dfs(tmp,ind+1)
for per in permutations(S):
  dfs("",0)
print(-1)

これを提出するとなんと3TLEになりました。
ざっとコードを見直しても計算量の見積もりが間違っているわけではなさそうです。
したがって、TLEの原因はPythonの再帰が遅いせいだと考えられます。
そこで、DFSを非再帰で書いてみます。
stackを使うとうまく書けます。（実装ではque=deque()としていますが、使い方としてはstackです）

import sys,math
from collections import defaultdict,deque
from itertools import combinations,permutations,accumulate,product
from bisect import bisect,bisect_left,bisect_right
from heapq import heappop,heappush,heapify
#from sortedcontainers import SortedList,SortedSet
def input(): return sys.stdin.readline().rstrip()
def ii(): return int(input())
def ms(): return map(int, input().split())
def li(): return list(map(int,input().split()))
#////////////////////////////
N,M = ms()
S = [input() for _ in range(N)]
T = set()
for _ in range(M):
  t = input()
  T.add(t)

def dfs():
  que = deque()
  que.append(("",0))
  while que:
    s,ind = que.pop()
    s += per[ind]
    if len(s)>16: continue
    if ind==N-1:
      if not s in T and 16>=len(s)>=3:
        print(s)
        exit()
      continue
    for i in range(1,16-len(s)+1):
      que.append((s+"_"*i,ind+1))
      
for per in permutations(S):
  dfs()
print(-1)

これを提出すると1732msでACすることができました！！
やはり横着して再帰のまま提出するのはよくないですね。
非再帰DFSを書いたことがない方はそこから勉強してみるとよいと思います。

おまけ

PyPyの文字列連結はO(N)なので、s+"_"*iのところが少し遅くなっていそうです。
そこで、グローバルなリストは定義し、それを使ってDFSをしてみます。
これはDFSで帰りがけに処理をするのと同じ方法で実装できます。

def dfs():
  que = deque()
  que.append((0,0))
  while que:
    global s
    ind,cnt = que.pop()
    if ind==-1:
      s.pop()
      s.pop()
      continue
    s.append(per[ind])
    cnt += len(per[ind])
    if cnt>16: continue
    if ind==N-1:
      tmp = ''.join(s)
      if not tmp in T and 16>=len(tmp)>=3:
        print(tmp)
        exit()
      continue
    for i in range(1,16-cnt+1):
      s += ["_"]
      que.append((-1,-1))
      que.append((ind+1,cnt+i))
      
for per in permutations(S):
  s = list()
  dfs()
print(-1)

これを提出すると1505msとなり、少し高速になりました。