かんプリンの学習記録

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

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つかうのもアリ