クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「2-1 迷路の最短路」(p.37)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「2-1 迷路の最短路」(p.37)
のコードをPythonで書き直したものとなっています。

入力

N\,M\\(大きさ{N}\times{M}の迷路)

入力例



10 10
#S######.#
......#..#
.#.##.##.#
.#........
##.##.####
....#....#
.#######.#
....#.....
.####.###.
....#...G#

 

解答

# 両端キューをインポートする
from collections import deque

INF = 100000000

# 入力
N, M = map(int,input().split())
maze = [list(input()) for _ in range(N)] # 迷路を表す文字列の配列
for x in range(N):
    for y in range(M):
        if maze[x][y] == 'S':            # スタートの座標
            sx = x
            sy = y
        if maze[x][y] == 'G':            # ゴールの座標
            gx = x
            gy = y
d = [[INF] * M for _ in range(N)]          # 各点までの最短距離の配列(INFで初期化)

# 移動4方向のベクトル
dx = [1, 0, -1, 0]
dy = [0, 1, 0, -1]

# (sx, sy)から(gx,gy)への最短距離を求める
# 辿り着けないとINF
def bfs():
    que = deque()
    # スタート地点を両端キューに入れ、その点の距離を0にする
    que.append([sx, sy])
    d[sx][sy] = 0

    # 両端キューが空になるまでループ
    while len(que) != 0:
        # キューの先端を取り出す
        p = que.popleft()
        # 取り出してきた状態がゴールなら探索をやめる
        if p[0] == gx and p[1] == gy:
            break
        # 移動4方向をループ
        for i in range(4):
            # 移動した後の点を(nx, ny)とする
            nx = p[0] + dx[i]
            ny = p[1] + dy[i]

            # 移動が可能かの判定とすでに訪れたことがあるかの判定(d[nx][ny] != INFなら訪れたことがある)
            if 0 <= nx < N and 0 <= ny < M and maze[nx][ny] != '#' and d[nx][ny] == INF:
                # 移動できるならキューに入れ、その点の距離をpからの距離+1で確定する
                que.append([nx, ny])
                d[nx][ny] = d[p[0]][p[1]] + 1
    return d[gx][gy]

res = bfs()
print(res)

Pythonで理解する蟻本「2-1 部分和問題」(p.34)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「2-1 部分和問題」(p.34)
のコードをPythonで書き直したものとなっています。

入力

n\,k\\a_1\,…\,a_n

入力例



4 13
1 2 4 7

再帰回数の上限

Pythonでは再帰回数の上限が設定されています。
そのため、競技プログラミングにおいては再帰関数の上限を十分大きい値に設定し直す必要があります。(AtCoderでは上限回数を超えると「Runtime Error」になります)

再帰回数の上限は

import sys
sys.setrecursionlimit(4100000)

で設定し直すことができます。

このエラーは初見では原因が分かりづらく、再帰回数に上限があることを知らなければエラーの原因を調べるのに大幅に時間を溶かすことになります。

この機会にしっかりと覚えておきましょう。


解答

# 再帰回数の上限を設定する
import sys
sys.setrecursionlimit(4100000)

# 入力
n, k = map(int,input().split())
a = list(map(int,input().split()))

# iまででSUMを作って、残りi以降を調べる
def dfs(i, SUM):
    # n個決め終わったら、今までの和SUMがkと等しいかを返す
    if i == n:
        return SUM == k
    
    # a[i]を使わない場合
    if dfs(i + 1, SUM):
        return True
    
    # a[i]を使う場合
    if dfs(i + 1, SUM + a[i]):
        return True
    
    # a[i]を使う使わないに拘わらずkが作れないのでfalseを返す
    return False

if dfs(0, 0):
    print("Yes")
else:
    print("No")

Pythonで理解する蟻本「2-1 キュー」(p.32)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「2-1 キュー」(p.32)
のコードをPythonで書き直したものとなっています。

PythonC++の違い

スタックとは異なり、キューはPythonにも存在します。

一方でlist(リスト)やdeque(両端キュー)のメソッドを利用して、キューとして使うこともできます。


またC++においては、popはコンテナの先頭側の要素を削除するだけで、返り値はありません

一方でPythonにおいては、popはコンテナの先頭側の要素を削除し、その値を返します(返り値がある)。

そのため、蟻本ではスタックの先頭側の要素を出力した後にその値を削除していますが、この記事のコードでは先頭側の要素を削除する際に値を変数に受け取り、その値を出力しています。

(蟻本のコードは「出力→値の削除」の順番になっているが、この記事のコードでは「値の削除→出力」の順番になっている。)


キューのコード

Queue(キュー)を使ったコード

from queue import Queue

que = Queue()    # キューを用意する
que.put(1)       # [] → [1]
que.put(2)       # [1] → [1,2]
que.put(3)       # [1,2] → [1,2,3]
x = que.get()    # 一番下から取り除いてxに代入する [1,2,3] → [2,3]
print(x)         # 1
x = que.get()    # [2,3] → [3]
print(x)         # 2
x = que.get()    # [3] → []
print(x)         # 3

list(リスト)を使ったコード

s = []          # リストを用意する
s.append(1)     # [] → [1]
s.append(2)     # [1] → [1,2]
s.append(3)     # [1,2] → [1,2,3]
x = s.pop(0)    # 一番下から取り除いてxに代入する [1,2,3] → [2,3]
print(x)        # 1
x = s.pop(0)    # [2,3] → [3]
print(x)        # 2
x = s.pop(0)    # [3] → []
print(x)        # 3

deque(両端キュー)を使ったコード

from collections import deque

s = deque()        #両端キューを用意する
s.append(1)        # [] → [1]
s.append(2)        # [1] → [1,2]
s.append(3)        # [1,2] → [1,2,3]
x = s.popleft()    # 一番下から取り除いてxに代入する [1,2,3] → [2,3]
print(x)           # 1
x = s.popleft()    # [2,3] → [3]
print(x)           # 2
x = s.popleft()    # [3] → []
print(x)           # 3

どれを使えばよい?

Queue、list、dequeを使って同じ処理をするコードを書いてみましたが、速度の点で言えばdequeの方が圧倒的に速いです。

なので、キューを使いたいときにはdequeを利用することをお勧めします。

どのくらい速度が違うのかは、sabaさんがQiita記事で詳しく説明されているので、興味のある方は一度読んでみるとよいでしょう。

qiita.com

Pythonで理解する蟻本「2-1 スタック」(p.31)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「2-1 スタック」(p.31)
のコードをPythonで書き直したものとなっています。

PythonC++の違い


C++と違い、Pythonにはスタックというデータ構造はありません。

Pythonではlist(リスト)やdeque(両端キュー)のメソッドを利用することで、スタックとして使うことができます。


またC++においては、popはコンテナの末尾側の要素を削除するだけで、返り値はありません

一方でPythonにおいては、popはコンテナの末尾側の要素を削除し、その値を返します(返り値がある)。

そのため、蟻本ではスタックの末尾側の要素を出力した後にその値を削除していますが、この記事のコードでは末尾側の要素を削除する際に値を変数に受け取り、その値を出力しています。

(蟻本のコードは「出力→値の削除」の順番になっているが、この記事のコードでは「値の削除→出力」の順番になっている。)


スタックのコード

listを使ったコード

s = []         # リストを用意する
s.append(1)    # [] → [1]
s.append(2)    # [1] → [1,2]
s.append(3)    # [1,2] → [1,2,3]
x = s.pop()    # 一番上から取り除いてxに代入する [1,2,3] → [1,2]
print(x)       # 3
x = s.pop()    # [1,2] → [1]
print(x)       # 2
x = s.pop()    # [1] → []
print(x)       # 1

deque(両端キュー)を使ったコード

from collections import deque

s = deque()    # 両端キューを用意する
s.append(1)    # [] → [1]
s.append(2)    # [1] → [1,2]
s.append(3)    # [1,2] → [1,2,3]
x = s.pop()    # 一番上から取り除いてxに代入する [1,2,3] → [1,2]
print(x)       # 3
x = s.pop()    # [1,2] → [1]
print(x)       # 2
x = s.pop()    # [1] → []
print(x)       # 1

listとdequeのどちらを使えばよい?

listとdequeの二つを使って書いてみましたが、速度の点で言えばdequeの方が圧倒的に速いです。

なので、スタックを使いたいときには、listではなくdequeを利用することをお勧めします。

どのくらい速度が違うのかは、sabaさんがQiita記事で詳しく説明されているので、興味のある方は一度読んでみるとよいでしょう。

qiita.com

Pythonで理解する蟻本「2-1 再帰関数」(p.30)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「2-1 再帰関数」(p.30)
のコードをPythonで書き直したものとなっています。

階乗を求めるコード

def fact(n):
    if n == 0:
        return 1
    return n * fact(n - 1)

フィボナッチ数列を求めるコード

def fib(n):
    if n <= 1:
        return n
    return fib(n - 1) + fib(n - 2)

メモ探索(メモ化再帰)によってフィボナッチ数列を求めるコード

MAX_N = 10 ** 8    # MAX_N は十分に大きい数(関数の実行時に、MAX_N >= n となるような整数)
memo = [0] * (MAX_N + 1)

def fib(n):
    if n <= 1:
        return n
    if memo[n] != 0:
        return memo[n]
    memo[n] = fib(n - 1) + fib(n - 2)
    return memo[n]

Pythonで理解する蟻本「1-6 ハードルが上がった「くじびき」」(p.25)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「ハードルが上がった「くじびき」」(p.25)
のコードをPythonで書き直したものとなっています。

入力

n\,m\\k_1\,…\,k_n

入力例



3 10
1 3 5


解答

二分探索とO({n}^3\,log\,{n})アルゴリズム

# 入力
n, m = map(int,input().split())
k = list(map(int,input().split()))

def binary_search(x):
    # xの存在し得る範囲k[l], k[l + 1], ..., k[r - 1]
    l = 0
    r = n
    
    # 範囲に何も含まれなくなるまで繰り返す
    while r - l >= 1:
        i = (l + r) // 2
        if k[i] == x:
            return True # 見つかった
        elif k[i] < x:
            l = i + 1
        else:
            r = i
    # 見つからなかった
    return False

# 二分探索を行うためにソート
k.sort()

f = False

for a in range(n):
    for b in range(n):
        for c in range(n):
            # 最も内側のループの代わりに二分探索
            if binary_search(m - k[a] - k[b] - k[c]):
                f = True

if f:
    print("Yes")
else:
    print("No")

O({n}^2\,log\,{n})アルゴリズム

# 入力
n, m = map(int,input().split())
k = list(map(int,input().split()))

# 2つで作れる数を格納する配列
kk = [0] * (n * n)

def binary_search(x):
    # xの存在し得る範囲kk[l], kk[l + 1], ..., kk[r - 1]
    l = 0
    r = n * n
    
    # 範囲に何も含まれなくなるまで繰り返す
    while r - l >= 1:
        i = (l + r) // 2
        if kk[i] == x:
            return True # 見つかった
        elif kk[i] < x:
            l = i + 1
        else:
            r = i
    # 見つからなかった
    return False

# k[c] + k[d]の取り得る数を列挙
for c in range(n):
    for d in range(n):
        kk[c * n + d] = k[c] + k[d]

# 二分探索を行うためにソート
kk.sort()

f = False
for a in range(n):
    for b in range(n):
        # 内側の2つのループの代わりに二分探索
        if binary_search(m - k[a] - k[b]):
            f = True

if f:
    print("Yes")
else:
    print("No")

モジュールを使った解法

27ページに書かれているように、このbinary_search関数のような関数は標準で使うことができます。

Pythonではbisectモジュールを使うことで上記と同様の処理を行うことができます。

bisectモジュールを使ったサンプルコードを下に載せておきます。

二分探索とO({n}^3\,log\,{n})アルゴリズム

# 二分探索するためにbisect_leftをimportする
from bisect import bisect_left

# 入力
n, m = map(int,input().split())
k = list(map(int,input().split()))

# 二分探索を行うためにソート
k.sort()

f = False

for a in range(n):
    for b in range(n):
        for c in range(n):
            # 最も内側のループの代わりに二分探索
            i = bisect_left(k, m - k[a] - k[b] - k[c])
            if 0 <= i < n and k[i] == m - k[a] - k[b] - k[c]:
                f = True

if f:
    print("Yes")
else:
    print("No")

O({n}^2\,log\,{n})アルゴリズム

# 二分探索するためにbisect_leftをimportする
from bisect import bisect_left

# 入力
n, m = map(int,input().split())
k = list(map(int,input().split()))

# 2つで作れる数を格納する配列
kk = [0] * (n * n)

# k[c] + k[d]の取り得る数を列挙
for c in range(n):
    for d in range(n):
        kk[c * n + d] = k[c] + k[d]

# 二分探索を行うためにソート
kk.sort()

f = False
for a in range(n):
    for b in range(n):
        # 内側の2つのループの代わりに二分探索
        i = bisect_left(kk, m - k[a] - k[b])
        if 0 <= i < n * n and kk[i] == m - k[a] - k[b]:
            f = True

if f:
    print("Yes")
else:
    print("No")


bisectモジュールの使い方を勉強する際には、bisect_leftとbisect_rightの違いを理解し、その使い分けを意識して勉強するとよいでしょう。

Pythonで理解する蟻本「1-6 三角形」(p.21)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「1-6 三角形」(p.21)
のコードをPythonで書き直したものとなっています。

入力

n\\a_1\,…\,a_n

入力例



5
2 3 4 5 10

解答

# 入力
n = int(input())
a = list(map(int,input().split()))

ans = 0 # 答え

# 棒を重複して選ばないようi < j < kとなるようにしている
for i in range(n):
    for j in range(i + 1, n):
        for k in range(j + 1, n):
            LEN = a[i] + a[j] + a[k]        # 周長
            ma = max(a[i], max(a[j], a[k])) # 最も長い棒の長さ
            rest = LEN - ma                 # 他の2本の棒の長さの和
            
            if ma < rest:
                # 三角形が作れるので、答えを更新できれば更新
                ans = max(ans, LEN)

# 出力
print(ans)

別解

22ページに書かれているように、この問題はO({n}\,log\,{n})で解くことができます。

まず、 O({n}\,log\,{n})で解く解法について説明します。

解法

n = 5
a = [3, 4, 2, 10, 5]

を入力例として考えてみましょう。


まず、aを昇順にソートしてみます。(O({n}\,log\,{n}))

n = 5
a = [2, 3, 4, 5, 10]

少し考えやすくなりました。(以下aは昇順に並んでいるものとして考えます)

ここで、周長が最大となる三角形の辺をa_i, a_j, a_k \left(1\leq{i}<{j}<{k}\leq{n}\right) とおくと

i + 1 = j

となることが分かります。以下、証明を載せておきます。


【証明】

周長が最大となる三角形の辺をa_i,a_j,a_k\,(1\leq{i}<{j}<{k}\leq{n})とおいたとき、i + 1 = jとなることを示す。

周長が最大となる三角形を三角形Pとし、その三辺をa_i,a_j,a_k\,(1\leq{i}<{j}<{k}\leq{n})…① とおいたとき、i + 1 \neq j…② となる三角形が存在すると仮定する。

このとき①, ②より、i + 2 < j…③ となる。

よって③よりa_i < a_(i+1) < a_jとなる。


また、①より三角形の成立条件から

a_i + a_j > a_k …④

となる。


ここで、a_{i+1},a_j,a_kを三辺とする三角形の存在について考えると

①より

a_{i+1} + a_k > a_j

a_j + a_k > a_{i+1}

となる。

また、④より

\begin{eqnarray}
a_{i+1} + a_j &>& a_i + a_j \\
&>& a_k
\end{eqnarray}

となりa_{i+1},a_j,a_kを三辺とする三角形が存在することが分かる。この三角形をQとおく。


このとき、三角形Pの周長をPL, 三角形Qの周長をQLとおくと

PL = a_i + a_j + a_k

QL = a_{i+1} + a_j + a_k


このときa_i < a_{i+1}から

\begin{eqnarray}
PL &=& a_i + a_j + a_k \\
&<& a_{i+1} + a_j + a_k \\
&=& QL
\end{eqnarray}

となる。


このことは三角形Pの周長が最大であることに反する。よって矛盾。

よって、周長が最大となる三角形の辺をa_i,a_j,a_k\,(1\leq{i}<{j}<{k}\leq{n})とおいたとき、i + 1 = jとなる。■



周長が最大となる三角形の辺をa_i,a_j,a_k (1\leq{i}<{j}<{k}\leq{n})とおくとi + 1 = jとなることが分かったので

b_i = a_i + a_{i+1} (1\leq{i}\leq{n-1})

と定義される数列bを考えると、この問題は

b_i,a_j (1\leq{i}<{i+1}<{k}\leq{n})

a_j < b_i

となるb_i,a_jについてb_i + a_jの最大値を求めよ。(ただし条件を満たすb_i,a_jが存在しない際には0を答えとする。)

という問題に言い換えることができます。


このとき、各b_iについて条件を満たすa_jは二分探索によってO({n}\,log\,{n})で調べることができます。

よって全体の計算量はO({n}\,log\,{n})となり、全探索で解く場合よりも大幅に計算量を減らすことができます。


別解のコード

from bisect import bisect_left

# 入力
n = int(input())
a = list(map(int,input().split()))

ans = 0 # 答え

# aを昇順にソートする
a.sort()

b = [a[i] + a[i+1] for i in range(n-2)]

for i in range(n - 2):
    j = bisect_left(a, b[i])
    if i + 3 <= j:
        ans = b[i] + a[j - 1]

# 出力
print(ans)