aacord’s memo

abcを中心にpythonで解いた問題のメモ、整理をしています。緑になった。

abc 105 D - Candy Distribution (python)

連続する部分列→累積和を取る
累積和をとって、差を考える
余りの種類は m 個あるけど、数列の長さから高々 n 個に座圧できる。
基本テクの詰め合わせ。

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from collections import defaultdict

n,m = map(int,readline().split())
a = list(map(int,readline().split()))
c = [0]*(n+1)
for i in range(n):
  c[i+1] = a[i]+c[i]

d = defaultdict(int)
for i in c:
  d[i%m] += 1
  
ans = 0
for i in d.values():
  ans += (i-1)*i//2
print(ans)

abc 074 D - Restoring Road Network (python)

サンプルを使ってワーシャルフロイドした結果が、サンプルと一致していなければ答えは -1 となる。
そうでない場合は存在する道の和を出力しないといけないが、それは i から j に道が引かれているとき == min( d[ i ][ k ] + d[ k ][ j ] ) > d[ i ][ j ] で求められる。
numpy 使うと便利

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np

from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import floyd_warshall

n = int(readline())
f = []
A = []
chk = []
for i in range(n):
  ar = list(map(int,readline().split()))
  f.append(ar)
  for j,ai in enumerate(ar):
    if i == j:
      continue
    A.append((i,j,ai))
  
INF = 10**9+1
g = [[INF] * n for _ in range(n)]
for a,b,c in A:
    g[a][b] = c
    g[b][a] = c

g = csr_matrix(g)
d = floyd_warshall(csgraph=g)

t = f == d
if t.all():
  ans = 0
  for i in range(n):
    d[i,i] = INF

  for i in range(n):
    for j in range(i+1,n):
      opt = np.min(d[i] + d[j])
      if opt > d[i,j]:
        ans += d[i,j]
        
  print(int(ans))
else:
  print(-1)

早稲田大学プログラミングコンテスト2019 B - 10 puzzle

場合分けが無限にある。
条件がきつい順に
全部0の場合→0
5が含まれていない場合→ No
全部5以下の場合→1
min(h, w) > 1 なら
最大値 9 → 4
最大値 8 → 3
最大値 6,7 → 2
となり、
そして上記以外で、h or w = 1 のときは、5の左右の列で最大値ごとに計算する必要がある。それを全ての5で計算して最小値をだす。えぐい。

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

h,w = map(int,readline().split())
c = [list(map(int,input().split())) for i in range(h)]
b = list(itertools.chain.from_iterable(c))
a = set(itertools.chain.from_iterable(c))

if a == {0}:
  print('Yes 0')
  exit()
  
if not 5 in a:
  print('No')
  exit()
  
if max(a) == 5:
  print('Yes 1')
  exit()
  
if w == 1 or h == 1:
  g = []
  p = 0
  m = 0
  for i in b:
    if i == 5:
      if m == 9:
        p += 3
      elif m == 8:
        p += 2
      elif m > 5:
        p += 1
      c = {0}
      g.append(p)
      p = 0
    else:
      m = max(i,m)
  if m == 9:
    p += 3
  elif m == 8:
    p += 2
  elif m > 5:
    p += 1
  g.append(p)
  k  = 100
  for i in range(len(g)-1):
    kk = max(g[:i+1])+max(g[i+1:])
    ks = min(ks,kk)
  print('Yes {}'.format(ks+1))
  exit()
  
if 9 in a:
  print('Yes 4')
  exit()
  
if 8 in a:
  print('Yes 3')
  exit()
  
print('Yes 2')

いろはちゃんコンテスト Day1 I - リスのお仕事

久しぶりに拡張ダイクストラの問題
day2 にも拡張ダイクストラが出ていたりする。
ダイクストラといっても、距離は一切出てこなくて、休憩回数の最小値を求める問題。
heap で (累計の休憩回数、前の隙間、今いる頂点)を保管して、今いる頂点から次の頂点に行くときに、その辺の隙間 != 前の隙間なら休憩回数を +1 して heap に入れ、= ならそのまま heap に入れる。
個人的には、同じ辺を二回以上通らないための工夫に悩んだが、考えると、heap でループしているので、ある頂点 i に初めて来た時の休憩回数が 1 から i に行くときの最小の休憩回数といえるので、heap で一回取り出された頂点にはもう行かなくていいと分かったので解決した。

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

from heapq import heappush, heappop

N,M,K = map(int,readline().split())
ABC = [list(map(int,readline().split())) for _ in range(M)]

graph = [[] for _ in range(N)]
for a,b,c in ABC:
    graph[a-1].append((b-1,c))
    graph[b-1].append((a-1,c))

ans = -1
used = [False]*N
used[0] = True
q = [(0,0,0)]
while q:
  t,pre,now = heappop(q)
  if now == N-1:
    ans = t
    break
    # 移動
  for nex,d in graph[now]:
    if used[nex]: #一回到達した頂点は考えなくていい
      continue
    else:
      used[nex] = True
      if pre == d:
        heappush(q,(t,d,nex))
      else:
        heappush(q,(t+1,d,nex))

if ans < 0:       
  print(ans)
else:
  print(ans*K)

abc 137 F - Polynomial Construction (python)

結局、フェルマーの小定理から x = i のとき 1 で、それ以外の時に 0 となるような p-1 次式が作れてしまうというだけの問題。賢い。
あとは a[ i ] = 1 のときに毎回 1 - (x-i)^(p-1) の係数を二項定理で作るだけ。

p = int(input())
a = list(map(int, input().split()))

MAX = 3000
fact = [1]*(MAX+1)
for i in range(1, MAX+1):
    fact[i] = (fact[i-1]*i) % p

inv = [1]*(MAX+1)
for i in range(2, MAX+1):
    inv[i] = inv[p % i]*(p-p//i) % p

fact_inv = [1]*(MAX+1)
for i in range(1, MAX+1):
    fact_inv[i] = fact_inv[i-1] * inv[i] % p

def comb(n, k):
    if n < k:
        return 0
    return fact[n] * fact_inv[n-k] * fact_inv[k] % p

ans = [0]*p
for i,ai in enumerate(a):
  if ai == 1:
    ans[-1] += 1
    for j in range(p):
      ans[j] -= pow(-i,j,p)*comb(p-1,j)
      ans[j] %= p
print(*ans[::-1])

abc 152 F - Tree and Constraints (python)

一つでもpath上に黒がある個数を、総数ーpath上がすべて白く塗る個数と置き換えて包除原理を使う。
m 個の path の中から i 個を選ぶ方法は itertools.combinations を使い、1つの path が通る辺の集合を辺1つずつに番号を座圧の要領でつけていき、2進数表記として保管した。複数の path を選ぶ時の辺集合を path1 | path2 | ... として表現することが出来る。

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from collections import deque,defaultdict

def all_bit_count(n):
    c = (n & 0x5555555555555555) + ((n>>1) & 0x5555555555555555)
    c = (c & 0x3333333333333333) + ((c>>2) & 0x3333333333333333)
    c = (c & 0x0f0f0f0f0f0f0f0f) + ((c>>4) & 0x0f0f0f0f0f0f0f0f)
    c = (c & 0x00ff00ff00ff00ff) + ((c>>8) & 0x00ff00ff00ff00ff)
    c = (c & 0x0000ffff0000ffff) + ((c>>16) & 0x0000ffff0000ffff)
    c = (c & 0x00000000ffffffff) + ((c>>32) & 0x00000000ffffffff)
    return c
  
n = int(readline())
d = defaultdict(int)
cnt = 1

g = [[] * n for i in range(n)]
for i in range(n-1):
  u, v = map(int,readline().split())
  u -= 1
  v -= 1
  g[u].append(v)
  g[v].append(u)

m = int(readline())
c = []
for i in range(m):
  u, v = map(int,readline().split())
  u -= 1
  v -= 1  
  que = deque()
  que.append((u,0))
  chk = [False]*n
  chk[u] = True
  while que:
      now,L = que.popleft()
      if now == v:
        c.append(L)
        break
      for i in g[now]:
          if not chk[i]:
            chk[i] = True
            mi = min(now,i)
            ma = max(now,i)
            mm = mi*100+ma
            if d[mm] == 0:
              d[mm] = cnt
              cnt += 1
            que.append((i,L|(1<<d[mm])))
                     
import itertools
ans = 1<<(n-1)
for i in range(1,m+1):
  for j in itertools.combinations(c,i):
    K = 0
    for k in j:
      K |= k
    l = all_bit_count(K)
    ans += pow(-1,i)*pow(2,n-1-l)
    
print(ans)

M-SOLUTIONS プロコンオープン D - Maximum Sum of Minimum

c を降順ソートすると、c [ i ] を新しく頂点に書き込むとき、スコアが c [ i ] になる回数は高々 i 回で、それまでに一つの辺でスコアを c [ i ] 以上獲得している回数は高々 i - 1 回あるので、求める最大値は sum( c [ 1 : ] ) になります。
そして、その構築方法は c の降順に dfs をして、スタートから i 回目にたどり着いた頂点に c [ i ] を書き込めば実現できます。

n = int(input())
g = [[] * n for i in range(n)]
for i in range(n-1):
  u, v = map(int, input().split())
  u -= 1
  v -= 1
  g[u].append(v)
  g[v].append(u)
  
d = list(map(int,input().split()))
d.sort(reverse=True)

from collections import deque
 
que = deque()
que.append(0)
chk = [False]*n
chk[0] = True
ans = [0]*n
ans[0] = d[0]
cnt = 1
while que:
  now = que.popleft()
  if cnt == n:
    break
  for i in g[now]:
    if not chk[i]:
      chk[i] = True
      ans[i] = d[cnt]
      cnt += 1
      que.append(i)
      
print(sum(d[1:]) if n > 1 else 0)
print(*ans)