かんプリンの学習記録

勉強したことについてメモしています. 主に競技プログラミングの問題の解説やってます.

ABC204のPythonによる実装

Pythonを使っている人に需要がありそうなので書きます.
※この記事は解法ではなく実装の解説を行います.解法が知りたい場合は別の記事を読んでください.

E以降は需要ありそうだったら書きます.

A問題

入力を受け取って計算結果を返すだけです.

# 入力
x,y = map(int,input().split())

# xとyが等しければ同じものを出力
if x == y:
    print(x)
# 異なるならxでもyでもないものを出力
else:
    print(3-x-y)

map(int,input().split())の部分は,入力された文字列input()split()によりスペース区切りで分けてmap(int,*)でそれぞれ整数に変換しています.

B問題

これも入力を受け取って計算結果を返します.

# 入力
N = int(input())
A = list(map(int,input().split()))

# Aの各要素について収穫する木の実の数を計算
ans = 0
for a in A:
    if a > 10:
      ans += a - 10

# 出力
print(ans)

for文の中身はこのような感じでもいいです.

for i in range(N):
    if A[i] > 10:
      ans += A[i] - 10

また,Aの各要素aについて収穫する木の実の数はmax(0,a-10)と表されるので

ans = sum(max(0,a-10) for a in A)

のように書くこともできます.

C問題

グラフの問題では,よく隣接リストを用います.隣接行列に比べて使用メモリが少なく済みます.

以下のプログラムではDFSを再帰で書いています.Pythonでは再帰はかなり時間がかかるので別の方法で実装します.
※以下のプログラムをそのまま提出してもREになります.これは再帰回数の上限が決まっているためです.上限を変更すれば通ります.

# 入力
N,M = map(int,input().split())
G = [[] for _ in range(N)] # 隣接リスト
for _ in range(M):
    A,B = map(int,input().split())
    G[A-1].append(B-1)

# 深さ優先探索(DFS)
def dfs(v):
    visited[v] = True
    cnt = 1
    for u in G[v]:
        if visited[u]:continue
        cnt += dfs(u,v)
    return cnt

ans = 0
# 全ての頂点でDFSを行う
for i in range(N):
    visited = [False] * N
    ans += dfs(i)
print(ans)

DFSを再帰ではなくstackを用いて行います.

# 入力
N,M = map(int,input().split())
G = [[] for _ in range(N)] # 隣接リスト
for _ in range(M):
    A,B = map(int,input().split())
    G[A-1].append(B-1)

# 深さ優先探索(DFS)
def dfs(s):
    visited = [False] * N
    visited[s] = True
    stack = [s] # 末尾に追加したりするstack
    while len(stack) > 0: # while stack: と書いてもいい
        v = stack.pop()
        for u in G[v]:
            if visited[u]: continue
            visited[u] = True
            stack.append(u)
    return sum(visited) # Trueの数

ans = 0
# 全ての頂点でDFSを行う
for i in range(N):
    ans += dfs(i)
print(ans)

おまけとして,さらに簡単な方法を紹介します.今回の制約だと隣接行列を作ることができるので以下の様なプログラムでも解けます.

import numpy as np
from scipy.sparse.csgraph import shortest_path

# 入力
N,M = map(int,input().split())
G = np.zeros((N,N)) # 隣接行列
for _ in range(M):
    A,B = map(int,input().split())
    G[A-1,B-1] = 1

# 到達可能な頂点の和を求める
ans = np.sum(shortest_path(G)!=np.inf)
print(ans)

D問題

公式解説のとおりに実装すると以下の様になります.

# 入力
N = int(input())
T = list(map(int,input().split()))

S = sum(T)
# dp[i][j] := T1~Tiからいくつか選んだ和をjにすることができるか
dp = [[False]*(S+1) for _ in range(N+1)]
dp[0][0] = True

for i in range(N):
    for j in range(S+1):
        dp[i+1][j] |= dp[i][j]
        if j+T[i] <= S:
            dp[i+1][j+T[i]] |= dp[i][j]

# S/2以上で最小のものを求める
for i in range((S+1)//2,S+1):
    if dp[N][i]:
        print(i)
        break

これを提出してもTLEになります.
ここでNumpyを用います.Numpyは多次元配列の演算を高速に行ったりしてくれるやつです.
Numpyを使うためにDP配列を使いまわすことで遷移の式を単純にしています.

import numpy as np
 
# 入力
N = int(input())
T = list(map(int,input().split()))
 
S = sum(T)
# dp[j] := いくつか選んで和をjにすることができるか
dp = np.zeros(S+1,np.bool_)
dp[0] = True
 
for i in range(N):
    dp[T[i]:] |= dp[:S-T[i]+1]

# S/2以上で最小のものを求める
print(np.argmax(dp[(S+1)//2:])+(S+1)//2)

感想

pypyつかうのもアリ

ABC204 D - Cooking

問題はこちら

問題概要

{N}個の料理があり,それぞれ作るためにオーブンを使う時間は{T_i}である.{2}つのオーブンを用いてすべての料理を作るのにかかる時間は最短何分か.

{1\leq N\leq 100\\
0\leq T_i\leq 10^3}

解説

{1}つのオーブンのみを用いることを考えると,{\sum_{i=1}^{N}T_i}が答え.

{2}つのオーブンだと,料理を{2}つのグループ{A,B}に分けた時の{\max(\sum_{i\in A}T_i,\sum_{i\in B}T_i)}がかかる時間となる.以降どのように料理を{2}つのグループに分けるかが問題となる.

{2^N}通り試していたら時間がかかるので別の方法を考える.

{S}{T_i}の総和({\sum_{i=0}^{N}T_i})とおく.
{A}を決めると{B}が決まるので{\max(\sum_{i\in A}T_i,\sum_{i\in B}T_i)=\max(\sum_{i\in A}T_i,S-\sum_{i\in A}T_i)}となるから,{\max(\sum_{i\in A}T_i,S-\sum_{i\in A}T_i)}を最小にする{A}はどのようなものかを考える.

ここで,{\sum_{i\in A}T_i\leq S-\sum_{i\in A}T_i}とする({\sum_{i\in A}T_i\geq S-\sum_{i\in A}T_i}のときは,{A}{B}を入れ替えることで{\sum_{i\in A}T_i\leq S-\sum_{i\in A}T_i}となる).このとき,かかる時間は{S-\sum_{i\in A}T_i}であり,これを小さくするには{\sum_{i\in A}T_i}を大きくする必要がある.{\sum_{i\in A}T_i\leq S-\sum_{i\in A}T_i}を変形すると{\sum_{i\in A}T_i\leq\frac{S}{2}}となるので,{N}個の料理をいくつか選んでできる{T_i}の総和のうち{\frac{S}{2}}以下の最大値が答えになる.これはナップサック問題とか部分和問題のようにDPで解ける.

ナップサック問題とか部分和問題を知らない人はEDPCとかで練習してください.

おまけ

慣れてないと{T_i}が大きい方から{\sum T}が少ないオーブンに振り分けていくような貪欲が思いつくかもしれませんが,様々な例を考えてみるとWAとなるようなものが見つかってしまいます(4,3,3,2,2など).このような直感に頼った貪欲は危険なのでやめましょう.面倒かもしれませんが貪欲の正当性の証明をしてください.

おまけ2

途中で最大値の最小化みたいな話が出てきますが,よく「最大値の最小化は二分探索!」と言われているので二分探索を考えてしまうかもしれません.ここで,なぜ最大値の最小化を二分探索により解くことができる問題があるかを説明します.

最大値の最小化は{min(max(A_1,A_2,\cdots))}のような形になり,{max(A_1,A_2,\cdots)}のような部分は扱うのが難しいことがあります.ここで,{max(A_1,A_2,\cdots)\leq k}であることと,任意の{i}に対して{A_i\leq k}であることは同値なので,{k}を決めうつことで{max(A_1,A_2,\cdots)}という{A_i}がたがいに関係し合う式から,{A_i\leq k}という各{i}について独立に考えることができる式に変換できます.さらに,「{i}に対して{A_i\leq k}であるか」の真偽は単調であるので二分探索が可能になります.

今回の問題では,{\min(\max(\sum_{i\in A}T_i,\sum_{i\in B}T_i))}です.{k}を決めたとき,{\sum_{i\in A}T_i\leq k}{\sum_{i\in B}T_i\leq k}は元の問題より簡単に求まるわけではなさそうなので二分探索で解くのは微妙です(解けないことはないが意味ない).

二分探索で解けるかな?と思うことは悪いことではないですが,このように考えることですぐに無理そうだな(意味ないな)と思うことができます.

例題:https://yukicoder.me/problems/no/710

提出プログラム

pythonのプログラムです.

import numpy as np
 
N = int(input())
T = list(map(int, input().split()))
S = np.sum(T)
DP = np.zeros(S//2+1,dtype=np.bool)
DP[-1] = 1
for i in range(N):
    if S//2+1-T[i] < 0: continue
    DP[:S//2+1-T[i]] |= DP[T[i]:] #numpyを用いていっぺんに遷移
print(S-S//2+np.argmax(DP))

感想

DPを学ぶとき一番最初にみるナップサック問題とか部分和問題に類似する問題でした.ちゃんと勉強しているかどうかが試されました.

ABC203 E - White Pawn

問題はこちら

問題概要

{(2N+1)×(2N+1)}のチェス盤があり,{(0,N)}に白のポーンがある.他にも{M}個の黒のポーンが置かれている.白のポーンは{(2N,*)}の方向に進むとき,白のポーンが{(2N,Y)}に到達できるような{Y}の数はいくつか.

{1\leq N\leq 10^9\\
0\leq M\leq 2×10^5}

解説

{dp[i][j]}を「{(i,j)}に到達できるか」とするようなDPを考える.遷移は

  • {dp[0][N]=True}
  • {dp[0][j]=False} ({j\neq N})
  • {dp[i][j]=dp[i-1][j-1]\lor dp[i-1][j+1]} ({(i,j)}に黒のポーンがある時)
  • {dp[i][j]=dp[i-1][j]} ({(i,j)}に黒のポーンがない時)

となる.このままだと{O(N^2)}とかになるが,黒のポーンがある時の遷移を行う回数は{O(M)}回であり,他の{dp[i][j]}の遷移はDP配列を使いまわすことで無視できる(いわゆるインラインDP).DP配列の長さが{O(N)}になるが,白のポーンが到達可能な範囲は{O(\min(N,M))}なので削れる.ソートがネックとなり{O(M\log M)}

提出プログラム

https://atcoder.jp/contests/abc203/submissions/23081411

感想

DよりEの方が初心者にとっては簡単?

yukicoder No.1518 Simple Combinatorics

問題はこちら

問題概要

{1}以上{N}以下の整数をランダムに選ぶという操作を{K}回行うとき,選ばれた整数の種類数の期待値に{N^K}をかけた値を求めよ.

{1\leq N,K < 10^9+7}

解説

種類数を{X}とする.
求めるものは{E[X]=\sum_{k=1}^{N}kP(X=k)}で, {P(X=k)}は包除原理とかで求められそうだが面倒なので別の方法を考える.

{\downarrow}

新たに以下の確率変数{Y}を導入する.

{\begin{equation}
Y_i= \left \{
\begin{array}{l}
1 (K回の操作で1度でもiを選んだ) \\
0 (K回の操作で1度もiを選ばなかった)
\end{array}
\right.
\end{equation}}

これを用いると,
{\displaystyle X=\sum_{i=1}^{N}Y_i}
となる. よって期待値の線形性より,
{\displaystyle E[X]=E[\sum_{i=1}^{N}Y_i]=\sum_{i=1}^{N}E[Y_i]}
したがって, 各{i}について{E[Y_i]=P(Y_i=1)}, つまり{K}回の操作で{1}度でも{i}を選ぶ確率を求めればいい.これはどの{i}{\displaystyle 1-\left(\frac{N-1}{N}\right)^K}なので答えは{N(N^K-(N-1)^K)}となる.

提出プログラム

https://yukicoder.me/submissions/661652

感想

ABC201 F - Insertion Sort

問題はこちら

問題概要

{N}がいて,左から{i}番目は人{P_i}である.以下の{3}種類の操作を繰り返し,人を昇順に並び替えるときのコストの最小値を求めよ.

  • コスト{A_i}を払い,人{i}を好きな位置に移動させる.
  • コスト{B_i}を払い,人{i}を左端に移動させる.
  • コスト{C_i}を払い,人{i}を右端に移動させる.

{1\leq N\leq 2×10^5\\
1\leq P_i\leq N\\
1\leq A_i,B_i,C_i\leq 10^9}

解説

{B_i\leftarrow \min(A_i,B_i),C_i\leftarrow \min(A_i,C_i)}とする.
それぞれの操作は元の位置によらないので,各人に対しての操作は高々{1}回でいい.

操作をしない人を決める(ただし番号は昇順)ことにより,他の人に対する操作が決まる.例として以下の問題を考える.

f:id:kanpurin:20210527012714p:plain:w400

{2},人{4},人{5}を操作しないと決める(選ばれなかった人は必ず動かす)

f:id:kanpurin:20210527015240p:plain:w400

{1}は左端に移動する操作,人{3}は好きな位置に移動する操作,人{6}は右端に移動する操作を行う.より一般的には,操作をしないと決めたすべての人よりも数が小さい人は左端に移動し,操作をしないと決めたすべての人よりも数が大きい人は右端に移動し,その他の人は好きな位置に移動する操作を行う必要がある.

番号の昇順に操作をしない人を選ぶように考えることで以下のようなDPで解ける.
{dp[i]を}{i}番目まで決めて最後に人{i}を選んだときの最小コスト」とする.元の順番で人{i}より左側にいるような人{j(j < i)}について,人{j}を選び,{j < x < i}となる人{x}は選ばないとすると{\displaystyle dp[j]+\sum_{k=j+1}^{i-1}A_k}の最小値が{dp[i]}となる.ただし,人{i}を最初に選ぶときは{\displaystyle\sum_{k=0}^{i-1}B_k}となる.
まとめると,{\displaystyle dp[i]=\min\left(\underset{j\in S}{\min}\left(dp[j]+\sum_{k=j+1}^{i-1}A_k\right),\sum_{k=0}^{i-1}B_k\right)}.ただし,{S}{j < i}かつ元の順番で人{j}が人{i}より左にいるような{j}の集合とする.

{\begin{eqnarray}
\displaystyle dp[i] &=& \min\left(\underset{j\in S}{\min}\left(dp[j]+\sum_{k=j+1}^{i-1}A_k\right),\sum_{k=0}^{i-1}B_k\right)\\
\displaystyle &=& \min\left(\underset{j\in S}{\min}\left(dp[j]-\sum_{k=1}^{j}A_k\right)+\sum_{k=1}^{i-1}A_k,\sum_{k=0}^{i-1}B_k\right)
\end{eqnarray}}

{\displaystyle\underset{j\in S}{\min}\left(dp[j]-\sum_{k=1}^{j}A_k\right)}が高速に求められれば解ける.

{i}の昇順に{\displaystyle\underset{j < i\land p(j) < p(i)}{\min}f(j)}みたいなものを求めるのは典型で,{i}の昇順にSegmentTreeの{p(i)}の位置に{f(i)}を入れていけば区間最小値を求める問題になって高速に解ける.

{dp[i]}が求まれば,人{i}の後の人は選ばないと考えることで{\displaystyle dp[i]+\sum_{j=i+1}^{N}C_j}の最小値が答えとなる.

提出プログラム

https://atcoder.jp/contests/abc201/submissions/22929907

感想

操作を行わない人を固定すると他の人の操作が決まることが分かれば後はすぐ解ける.試行錯誤して見つけていくしかない?

ABC202 E - Count Descendants

問題はこちら

問題概要

{N}頂点の根付き木について以下の{Q}個のクエリに答えよ.

整数{U,D}が与えられる.次の条件を満たす頂点{u}の個数を求めよ.

  • {u}から根への最短パス上に頂点{U}が存在する.
  • {u}から根への最短パスに含まれる辺の数がちょうど{D}である.

{2\leq N\leq 2×10^5\\
1\leq Q\leq 2×10^5}

解説

根から頂点{u}までの距離を{d_u}とします.求めるものは頂点{U}を根とする部分木内の{d_u=D}となる頂点の数です.

DFSで訪れた順で頂点を並べてみると,ある部分木の頂点は連続していることがわかります(典型).よってDFSで頂点を見ていくことにします.現在見た中で{x=d_u}となる頂点{u}の数を{C(x)}として,これを記録しながらDFSを行います.ある頂点{u}を訪れた時の{C(x)}の状態と去る時の{C(x)}の状態の差により{U=u}となるクエリの答えが求められます.保持しておく「頂点{u}を訪れた時の{C(x)}の状態」は{U=u}となるクエリの{C(x)}の値だけでいいので{O(N+Q)}で解けました.

あとからもっと詳しく書くかも.

これとかこれっぽいですね.←ネタバレ注意

提出プログラム

https://atcoder.jp/contests/abc202/submissions/22838907

感想

DFSしながら状態を更新していくのは良く出ますね.

yukicoder No.1506 Unbalanced Pocky Game

問題はこちら

問題概要

長さ{N}の数列{A}があり,{A}の末尾の要素を{k}として,末尾の要素を{0}以上{k}未満の整数に変えて{0}になった要素は消すという操作を{2}人のプレーヤーが交互に行う.操作ができなくなる(自分の操作開始時に{A}が空)と負けとなるとき,どちらのプレーヤーが勝つか.

{1\leq N\leq 2×10^5\\
1\leq A_i\leq 10^9}

解説

grundy数を考える.{grundy_i(x)}を「{|A|=i}かつ{A_i=x}」のgrundy数とすると

{grundy_i(x)=\left\{\begin{array}{ll}
0 & (i=1\land x=0)\\
grundy_{i-1}(A_{i-1}) & (2\leq i \land x=0)\\
x-1 & (1\leq x\leq grundy_i(0))\\
x & (grundy_i(0) < x)
\end{array}
\right.}

であり,{grundy_N(A_N)=0}なら後手のプレーヤー,{grundy_N(A_N)\neq 0}なら先手のプレーヤーの勝ち.

こうすれば,複数のゲームを並列に行う問題でも解ける.

提出プログラム

https://yukicoder.me/submissions/658273

感想