AtCoder Beginner Contest 149

*ずっと下書きに保存して投稿を忘れた。

---

昨日のAtCoder Grand Contest 041には、残念ですが用事があって参加できなかった。
が、現時点自分のレベルで良い成績が取れるとは思えない。

今回自分は4問しか解けなかったが、順位表を見ると300位以降5.6問を解けた人が殆どない。
つまり4と5の間に大きいギャップがあると思う。

###C - Next Prime

ある数が素数であるかについて、
もしこの数がすべて自分の平方根より小さい数で整除されないなら、素数のはずです。

import sys
readline = sys.stdin.buffer.readline

x=int(readline())


def prime(n):
  for i in range(2, int(n ** 0.5 + 1)):
    if n % i == 0:
      return False
  return True
  

while True:
  if prime(x):
    print(x)
    quit()
  x += 1

###D - Prediction and Restriction

この問題を見て、とある会社「S*****x」のプログラミング選考が思い出された。
Pythonの辞書を使うと、便利にじゃんけんで出す手を決める。
というわけで自分が３つ辞書を作った。

この問題で注意すべきな点がただ一つ：
今回筐体の出す手がK回前と同じの時、
K回後の勝利の為、今回は自在に「あいこ」か「負け」を選ぶべき。

import sys
readline = sys.stdin.buffer.readline

n,k=map(int,readline().split())
r,s,p=map(int,readline().split())
t=input()
dic={'r':p,'s':r,'p':s}  # スコア
dic2={'r':'p','s':'r','p':'s'}  # 勝つ
dic3={'r':'s','s':'p','p':'r'}  # 負け
h='' 　#毎回自分の出す手を記録する

res=0
for i in range(k):  # K回目までの点数をまず計算する
  res+=dic[t[i]]
  h+=dic2[t[i]]

for i in range(k,n):
  if h[i-k]!=dic2[t[i]]:  # 普通な場合に当然「勝つ」を選ぶ
      res+=dic[t[i]]
      h+=dic2[t[i]]
  else:
    if i+k<n:  # K回後ゲームが終わっていない場合を考える
      if dic2[t[i+k]]!=t[i]:  # 「あいこ」できる場合
        h+=t[i]
      else:
        h+=dic3[t[i]]  # 「負け」を選ぶ
    else:
      h+=t[i]  # 「あいこ」か「負け」のどっちでも同じ
print(res)

###E - Handshake

maspyさんのコードを参考いたします。
https://atcoder.jp/contests/abc149/submissions/9219994

まずはデータを読む。

  
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline

# 価値x以上の握手を全て行うとする
# 握手回数が決まる。左手ごとに集計できる。

import numpy as np

N,M = map(int,readline().split())
A = np.array(read().split(),np.int64)

A.sort()

次に「価値x以上の握手」を数える。

def shake_cnt(x):
    # 左手がA[i]の時、価値x以上になるため選べない右手を数える
    # x-A[i]がmax(A)に超えると挿入点が一番右になって、すべての手が選べないことがわかる
    X = np.searchsorted(A,x-A)
    return N * N - X.sum()

np.searchsorted(list, x)はPython基礎libのbisect.bisect_leftと同じ、
ソートされた順序を保ったままxをlistに挿入できる点を探し当てます。
当然、numpyがarrayの場合でも処理できます。

arr = np.array([1,2,3,5])
print(np.searchsorted(arr, [0,3,4,6]))
# [0 2 3 4]

import bisect
bisect.bisect_left([1,2,3,5], 4)
# 3
bisect.bisect_left([1,2,3,5], [0,3,4,6])
# TypeError

後は二分探索でのX探しと幸福値計算。

left = 0  # 握手の回数がM以上
right = 10 ** 6  # 握手の回数がM未満
while left + 1 < right:
    x = (left + right) // 2
    if shake_cnt(x) >= M:
        left = x
    else:
        right = x

'''
例3
（9 14
1 3 5 110 24 21 34 5 3）
の場合、
left=113,right=114,shake_cnt(left)=15,shake_cnt(right)=11
'''

X = np.searchsorted(A,right-A)  # 行わない人数
shake = N * N - X.sum()
Acum = np.zeros(N+1,np.int64)
Acum[1:] = np.cumsum(A)  # 累積和

happy = (Acum[-1] - Acum[X]).sum() + \  # 右手のパワー:
        (A * (N - X)).sum()  # 左手のパワー
happy += (M - shake) *left  # 残り(14-11=)3つ握手の価値=left=113
print(happy)

おまけの高速フーリエ変換：
https://atcoder.jp/contests/abc149/submissions/9228535

###F - Surrounded Nodes

import sys
readline = sys.stdin.readline

N = int(readline())
graph = [[] for _ in range(N + 1)]
for i in range(N - 1):
    a, b = map(int, readline().split())
    graph[a].append(b)
    graph[b].append(a)

parent = [0] * (N + 1)
order = []
stack = [1]
while stack:
    x = stack.pop()
    order.append(x)  # 行きがけ順探索リスト
    for y in graph[x]:
        if y == parent[x]:
            continue
        parent[y] = x  # 親ノードを記録
        stack.append(y)

MOD = 10**9 + 7
half = pow(2, MOD - 2, MOD)
power_inv = [1] * (N + 1)
size = [1] * (N + 1)
for i, v in enumerate(order[::-1], 1):
    p = parent[v]
    x = size[v]  # vの子孫ノード数（自分も含む）をとる
    size[p] += x  # 親にノード数を加算
    power_inv[i] = power_inv[i - 1] * half % MOD  # [1, 1/2, 1/4, ...]

ans = sum((1 - power_inv[i] - power_inv[N - i] + power_inv[N]) %
          MOD for i in size[2:])  # 辺iがSに含まれるための必要十分条件による辺の個数の期待値
ans += 1  # 木の頂点の個数は辺の個数+1
ans -= power_inv[N] + N * power_inv[1]  # -「すべての辺が含まれない、つまり空グラフの場合の誤加算」-N/2
ans %= MOD

print(ans)