ラーメン、それは人類が到達した光

日々の気づきと思ったことをそのまま書き連ねるブログ

Pythonで解く蟻本(2-2: 猪突猛進! "貪欲法")

※自分の勉強の備忘録として書くので、ちゃんとした説明はない記事です。
記事の構成としては、ざっくりとこんな感じです。

  1. 蟻本の答案をPython3に変換したもの
  2. AtCoderの例題のPython3解答

蟻本と似たような問題は、Qiitaの記事(AtCoder 版!蟻本 (初級編) - Qiita)を参考に探しました。
AtCoderの解答解説を読んでもコーディングで躓く、他の人のACコードを見ても何が起きてるか分からない、という方の助けになれば幸いです。

この記事では、貪欲法の話しか出てこないです。
なお、持っている蟻本は以下の第2版です。

硬貨の問題(蟻本 P.42)

2-1までの深さ優先探索や幅優先探索のコードと比較するとかなりシンプルですね。
(実際の問題だと、そもそも法則を見つけにくかったりするんだろうな...)

できるだけ高い金額のコインから消費していきます。

V = [1, 5, 10, 50, 100, 500]
C = list(map(int, input().split()))

A = int(input())

ans = 0
for i in reversed(range(6)):
	t = min(A // V[i], C[i])
	A -= t * V[i]
	ans += t

print(ans)

JOI 2007 予選 A おつり

A - おつり

これは上記の問題とほぼ同じ感じで解けますね。
できるだけ高い金額のコインから消費するものの、コインの上限がなくなったパターンです。

values = [1, 5, 10, 50, 100, 500]
buy = int(input())
ans = 0

rest = 1000 - buy

for i in reversed(range(6)):
    coin = rest // values[i]
    ans += coin
    rest -= coin * values[i]

print(ans)

AOJ Course コイン問題

Coin Changing Problem

一瞬同じ問題に見えつつも、先ほどまでと同様にやっていてはいきなり躓くことがサンプルケースの時点で分かります。

コインの金額の差が不揃いなのが原因のようですね。
nが5,000以下なので、二重ループでも2.5e+7の計算量となるので間に合うため、全探索します。

INF = 10**8

n, m = map(int, input().split())
c = list(map(int, input().split()))

d = [INF] * (n + 1)
d[0] = 0

for coin in c:
	for price in range(coin, n + 1):
		d[price] = min(d[price], d[price - coin] + 1)

print(d[n])

区間スケジューリング問題(蟻本 P.43)

これは、コーディング自体は全く難しくないですね。
終端に注目する、という所に気付けずにハマってしまうと辛そうな雰囲気の問題です。

N = int(input())
S = list(map(int, input().split()))
T = list(map(int, input().split()))

itv = [[0] * 2 for _ in range(N)]

for i in range(N):
	itv[i][0] = T[i]
	itv[i][1] = S[i]

itv.sort()

ans = 0
t = 0
for i in range(N):
	if(t < itv[i][1]):
		ans += 1
		t = itv[i][0]

print(ans)

キーエンス プログラミング コンテスト 2020 B - Robot Arms

B - Robot Arms

これも、上記とほぼ同じ形で解けるので簡単ですね。
ロボットアームの右側に着目します。
と言いつつ、この流れで解いてるから簡単なだけで実際には結構悩む気がしました。

こういう問題の解き方にパッと気付けるようになりたい。

n = int(input())
position = []

for i in range(n):
	x_in, l_in = map(int, input().split())
	position.append([x_in + l_in, x_in - l_in])

position.sort()
ans = 0
t = - (10 ** 9 + 1)

for i in range(n):
	if t <= position[i][1]:
		ans += 1
		t = position[i][0]

print(ans)

ABC 103 D - Islands War

D - Islands War

これも同じように実装すればいけますね。
島の番号が大きい方から橋を分断していきます。

n, m = map(int, input().split())
ba = []

for i in range(m):
    ba.append(list(map(int, input().split())))
    ba[i][0], ba[i][1] = ba[i][1] - 1, ba[i][0] - 1

ba.sort()
bridge = [ba[0][0] - 1]


for i in range(1, m):
    if ba[i][1] > bridge[-1]:
        bridge.append(ba[i][0] - 1)

print(len(bridge))

Codeforces 296 DIV1 B Clique Problem

Problem - B - Codeforces

ちょっと問題の意味が分からず、、断念しました。

ABC 038 D プレゼント

D - プレゼント

これは難しかったので解答を見てしまいました。
なるほど。片方を昇順に並べてもう片方を降順に並べる...!
思いつかなかった...!

bisect関数を二次元リストに使ってるのも個人的にはなるほどという感じでした。

from bisect import bisect_left
from operator import itemgetter

n = int(input())
box = [list(map(int, input().split())) for i in range(n)]

# h_i は昇順にした上で、h_i が等しい場合 w_i は降順にソートする
box = sorted(box, key=itemgetter(0), reverse=True)  # w_i
box = sorted(box, key=itemgetter(1))  # h_i

dp = [float("inf")] * n

for i in range(n):
    dp[bisect_left(dp, box[i][0])] = box[i][0]

print(bisect_left(dp, float("inf")))

Best Cow(蟻本 P.45)

C++ の for文と Python の for文 の仕様が違うのでwhile文で実装します。

n = int(input())
s = input()
t = ""

a, b = 0, n - 1
while a <= b:
  left = False
  i = 0
  while a + i <= b:
    if s[a + i] < s[b-i]:
      left = True
      break
    elif s[a + i] > s[b - i]:
      left = False
      break
    i += 1

  if left:
    t += s[a]
    a += 1
  else:
    t += s[b]
    b -= 1

print(t)

ABC 076 C Dubious Document 2

C - Dubious Document 2

これは、以下のような考えでコードを書いてみました。

  1. Sの中にTが含まれている場合は、?を全部aに変換する
  2. 含まれていない場合は、Sの後ろの文字からT文字数文抜き出してTと比較
  3. Tが見つかったら、SにTを埋め込んだあとに残った?を全部aに変換する

?をaに変換するときは、replace関数を使った方が良いことにはてな投稿時に気付きましたが気にしないことにします。

S = list(input())
T = list(input())
flag = False

for i in reversed(range(len(S) - len(T) + 1)):
  T_flg = S[i: i + len(T)]
  for j in range(len(T)):
    if T_flg[j] != T[j] and T_flg[j] != "?":
      break
  else:
    flag = True
    S[i:i + len(T)] = T
    break

if flag == False:
  print("UNRESTORABLE")
else:
  for i in range(len(S)):
    if S[i] == "?":
      S[i] = "a"
  print("".join(map(str,S)))

ABC 007 B 辞書式順序

B - 辞書式順序

これはめちゃくちゃ簡単ですね...!

A = input()

if A == "a":
  print(-1)
else:
  print("a")

ABC 009 C 辞書式順序ふたたび

C - 辞書式順序ふたたび

これは難しかったので、解説と他の人のAC解答をみました。
抽象的には理解したものの、辞書型変数同士の減算やCounter関数がちょっと独特でまだ頭の中でイメージしにくいのでこれは修行が必要な予感がします。

from collections import Counter

n, k = map(int, input().split())
s = list(input())

s_sort = sorted(s)
t = ""

# 元の位置と違う文字の数
diff = 0
# 見終わったけどまだ使ってない文字
counter = Counter(s[:1])
counts = sum(counter.values())

# 1文字目から順に見ていく
for i in range(n):
    # t + c が可能か調べる
    for c in s_sort:
        # t + c の中で元の位置と違う文字の数
        if c != s[i]:
            diff1 = diff + 1
        else:
            diff1 = diff

        # まだ使ってない文字の中で元の位置と違う文字の数
        if counter[c] > 0:
            diff2 = counts - 1
        else:
            diff2 = counts
        # 両方を足して k 以下ならば t + c が可能
        if diff1 + diff2 <= k:
            t += c
            s_sort.remove(c)
            diff = diff1
            counter = Counter(s[:i + 2]) - Counter(t)
            counts = sum(counter.values())
            break
print(t)

Saruman's Army (POJ No.3069)

これは結構直感的に解けますね。

n = int(input())
r = int(input())
x = list(map(int, input().split()))

x.sort()
i = 0
ans = 0

while i < n:
    s = x[i]
    i += 1
    while i < n and x[i] <= s + r:
        i += 1
    p = x[i - 1]

    while i < n and x[i] <= p + r:
        i += 1
    ans += 1

print(ans)

ABC 083 C Multiple Gift

C - Multiple Gift

数列Aが最大となるためには、常に2倍ずつ大きくするのが最良なので以下のコードになります。

x, y = map(int, input().split())
ans = 0

while x <= y:
    x *= 2
    ans += 1

print(ans)

ARC 006 C 積み重ね

C - 積み重ね

床に積まれた段ボールの山の中に、新しくトラックから運んできた段ボールが積める山がある場合は絶対に積む、というロジックで解決できます。

n = int(input())
w = [int(input()) for _ in range(n)]
floor = [w[0]]

for i in range(1, n):
    for j in range(len(floor)):
        if floor[j] >= w[i]:
            floor[j] = w[i]
            break
    else:
        floor.append(w[i])

print(len(floor))

ABC 005 C おいしいたこ焼きの売り方

C - おいしいたこ焼きの売り方

たこ焼きを古い順に確認して、お客さんが来る時間 - t 以降に作られたたこ焼きが今あるかどうか確認していきます。

t = int(input())
n = int(input())
a = list(map(int, input().split()))
m = int(input())
b = list(map(int, input().split()))

flag = True

if n < m:
    print("no")
    exit()

j = 0
ans = 0

for i in range(m):
    while j < n:
        if b[i] - t <= a[j] <= b[i]:
            j += 1
            ans += 1
            break
        else:
            j += 1

if ans == m:
    print("yes")
else:
    print("no")

SRM 560 DIV1 Easy TomekPhone

TopCoder Statistics - Problem Statement

自分自身の理解が浅いからなのか、TopCoderの問題文に辿り着けずどんな問題か分かりませんでした...!

ABC 091 C 2D Plane 2N Points

C - 2D Plane 2N Points

これは最初、マッチする点が少ない物同士をペアにしていけば良いかと思ったのですが、そうではなかったです。
解説をみたら納得したのですが、もっとシンプルに、青い点をx座標が小さい順に見ていって、赤い点をy座標が大きい順にペアにしていけば良かったようです。
うーん、やっぱり貪欲法はパッと考えると難しいですね。

from operator import itemgetter

n = int(input())
ab = [list(map(int, input().split())) for _ in range(n)]
cd = [list(map(int, input().split())) for _ in range(n)]
ans = 0
dp_ab = [False] * n
dp_cd = [False] * n

ab.sort(key = itemgetter(1), reverse = True)
cd.sort()

for i in range(n):
    for j in range(n):
        if cd[i][0] > ab[j][0] and cd[i][1] > ab[j][1] and dp_cd[i] == False and dp_ab[j] == False:
            dp_cd[i] = True
            dp_ab[j] = True
            ans += 1

print(ans)

Fence Repair (POJ No.3253)

だいぶ表現の仕方は違いますが、Pythonだったらこんな感じでよりすっきり書ける気がしています。

import bisect

N = int(input())
L = list(map(int, input().split()))

L.sort()
ans = 0

while N > 1:
    mii1 = 0
    mii2 = 1

    t = L[mii1] + L[mii2]
    ans += t
    L.pop(mii2)
    L.pop(mii1)
    idx = bisect.bisect_left(L, t)
    L.insert(idx, t)
    N -= 1

print(ans)

Codeforces 263 DIV2 C Appleman and Toastman

Problem - C - Codeforces

applemanが、常に一番小さい数字とそれ以外のペアに分割していくのが最適解となります。
そうすると、一番小さい数字は2回, 二番目に小さい数字は3回, ... n-1番に小さい数字はn回で、n番目に小さい(つまり一番大きい)数字もn回の足し算となります。

なので、以下のコードとなりますね。

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

a.sort()
score = sum(a) + a[n - 1] * (n - 1)

for i in range(1, n):
    score += a[i - 1] * i

print(score)

まとめ

うーん、貪欲法のまとめはなんていうか、最適なパターンを思いついたら、とりあえず考えてみて実装してみる!ということでしょうか...?

Pythonで解く蟻本(2-1: すべての基本 "全探索")

※自分の勉強の備忘録として書くので、ちゃんとした説明はない記事です。
記事の構成としては、ざっくりとこんな感じです。

  1. 蟻本の答案をPython3に変換したもの
  2. AtCoderの例題のPython3解答

蟻本と似たような問題は、Qiitaの記事(AtCoder 版!蟻本 (初級編) - Qiita)を参考に探しました。
AtCoderの解答解説を読んでもコーディングで躓く、他の人のACコードを見ても何が起きてるか分からない、という方の助けになれば幸いです。

この記事では、ほとんど深さ優先探索と幅優先探索の話しか出てこないです。
なお、持っている蟻本は以下の第2版です。

部分和問題(蟻本 P.34)

これは、記載されている解答をそのまま Python で書けばいいので簡単ですね。

def dfs(i: int, s: int):

  if i == n:

    return s == k



  if dfs(i + 1, s):

    return True



  if dfs(i + 1, s + a[i]):

    return True



  return False



n = int(input())

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

k = int(input())



if dfs(0, 0):

  print("Yes")



else:

  print("No")

ABC 045 C - たくさんの数式

C - Many Formulas

これは上記の蟻本の問題と同じパターンです。ただし、コーディングに少し工夫が必要でした。
pythonの文字列操作で解決できることに気付いたら早そうです。

s = input()

n = len(s)



def dfs(i: int, t: str):

  if i == n-1: #i = n - 1に到達したら終了

    # + で区切った数字をint型に変換して足し算

    return sum(list(map(int, t.split('+'))))



  # 文字列で次の数字をそのまま連結するパターンと、+で区切るパターン

  return dfs(i+1, t+s[i+1]) + dfs(i+1, t + '+' + s[i+1])



# 初期値でsの一文字目を入れておく

print(dfs(0, s[0]))

ABC 079 C - Train Ticket

C - Train Ticket

ちょっと汚いですが、こんな感じになりました。
正直、dfsについてarrayを引数としなくても良かった気がしています。

def dfs(array: list, index: int, calc: str, result: int):

    if index == len(array):

        if result == 7:

            print("{}=7".format(calc))

            exit(0)



    if index == 0:

        dfs(array, index + 1, array[index], int(array[index]))

    elif index < len(array):

        dfs(array, index + 1, calc + "-" + array[index], result - int(array[index]))

        dfs(array, index + 1, calc + "+" + array[index], result + int(array[index]))



array = list(input())

dfs(array, 0, "", 0)

上記のコードでは exit で処理を完了していますが、この代わりに return を使うとprintしすぎてダメになります。

exitとreturnの違いは、以下になります。

  • exitは呼び出された時点で全ての処理を即終了
  • returnは呼び出されている関数の処理を終了

要は、return print(~) としてしまうと、正解パターンを全部吐き出してしまうのでWAとなるわけです。
この記事が分かりやすかったので気になる人はどうぞ。

» exitとreturnの終了処理は何が違うのか?技術ブログ


ちなみに、return を使って実装した場合の一例はこちらになります。
ちょっと複雑になりますね。

def dfs(array: list, index: int, calc: str, result: int):

    if index == len(array):

        if result == 7:

            return "{}=7".format(calc)

        else:

            return ""



    if index == 0:

        return dfs(array, index + 1, array[index], int(array[index]))

    elif index < len(array):

        return dfs(array, index + 1, calc + "-" + array[index], result - int(array[index])) + dfs(array, index + 1, calc + "+" + array[index], result + int(array[index]))



array = list(input())

S = dfs(array, 0, "", 0)

print(S[:9])

ABC 104 C - All Green

C - All Green

これは解答を読んでしまいました。
AtCoderの解説を読むと、以下の2点を前提とするようです。

  • 中途半端に解く配点は 1 種類以下であり、それ以外の配点は完全に解くかまったく解かない。
  • 中途半端に解く配点があるなら、それは完全に解く配点以外の配点の中で最も高い配点である。

直感的にはなんとなく分かる気がしつつ、「本当にそうなの?完答の配点に依らない?」と思ってしまいます。
が、完答の配点によって最小解答数が変わってくるのは事実ですが、上記のパターンの中にしか最小解答数が存在しないことはよく考えればその通りでした。

うーん、悔しいですね。

def dfs(i, sum, count, nokori):

    global ans

    if i == d:

        # G 点に満たなければ nokori のうち一番大きいものを解く

        if sum < g:

            use = max(nokori)

            # 解く問題が問題数を超えないように注意

            n = min(pc[use - 1][0], -(-(g - sum) // (use * 100))) #ムダに-演算子を重ねているのは切り上げしたいから

            count += n

            sum += n * use * 100



        if sum >= g:

            ans = min(ans, count)

    else:

        # 総合スコア、解いた問題数、まだ解いてない問題を更新

        dfs(i + 1, sum, count, nokori)

        dfs(i + 1, sum + pc[i][0] * (i + 1) * 100 + pc[i][1], count + pc[i][0], nokori - {i + 1})





d, g = map(int, input().split())

pc = [list(map(int, input().split())) for i in range(d)]



ans = float("inf")



dfs(0, 0, 0, set(range(1, d + 1)))

print(ans)

コード中のコメントに記載してるのですが、Pythonで切り上げをする時のテクニック、面白いですね...!

Python で割り算をするときの切り上げと切り捨て | 民主主義に乾杯

ARC 029 A - 高橋君とお肉

A - 高橋君とお肉

これは、ここまでに出てきた中で一番素直ですね。
以下にコードを挙げます

def dfs(i: int, grill1: int, grill2: int):

    if i == n:

        return max(grill1, grill2)

    elif i < n:

        return min(dfs(i + 1, grill1 + t[i], grill2), dfs(i + 1, grill1, grill2 + t[i]))



n = int(input())

t = []



for i in range(n):

    t.append(int(input()))



print(dfs(0, 0, 0))

ABC 002 D - 派閥

D - 派閥

複雑性が増しました。
他人のACコードを参考にしてしまったので悔しいです。

itertools.combinations() 関数で組み合わせを列挙できるのが凄く便利。
(これを使ってる解答が多かった。)

ということで、以下が解答になります。

import itertools



n, m = map(int, input().split())

t = [[0]*n for _ in range(n)]



def dfs(i, group):

    global ans

    if i == n:

        flag = True

        for i in itertools.combinations(group, 2):

            if t[i[0]][i[1]] == 0:

                flag = False

                break

        if flag:

            ans = max(ans, len(group))

    else:

        dfs(i + 1, group)

        dfs(i + 1, group + [i])



for _ in range(m):

    x, y = map(int, input().split())

    t[x-1][y-1] = 1

    t[y-1][x-1] = 1



ans = 0

dfs(0, [])

print(ans)

Lake Counting(蟻本 P.35)

これは蟻本の解答を素直にPythonで書けば良いだけですね。

n, m = map(int, input().split())

field = ["" for _ in range(n)]



for i in range(n):

    field[i] = list(input())



def dfs(x, y):

    field[x][y] = "."



    for dx in range(-1, 2):

        for dy in range(-1, 2):

            nx = x + dx

            ny = y + dy



            if 0 <= nx < n and 0 <= ny < m and field[nx][ny] == "W":

                dfs(nx, ny)



    return



res = 0

for i in range(n):

    for j in range(m):

        if field[i][j] == "W":

            dfs(i, j)

            res += 1



print(res)

ATC 001 A 深さ優先探索

A - 深さ優先探索

自分で実装したのは以下のコードなのですが、一部のケースでTLEを吐いてしまいACできませんでした。
後述してますが、PyPyでコード提出していたのが原因でした。Pythonでコード提出すればACになります。

import sys

sys.setrecursionlimit(10 ** 8)



def dfs(i: int, j: int):

	global flg

	memo[i][j] = 1

	

	if c[i][j] == "g":

		flg = True

		print("Yes")

		exit(0)



	for a in range(4):

		if 0 <= i + dy[a] < h and 0 <= j + dx[a] < w:

			if c[ i+dy[a] ][ j+dx[a] ] != "#" and memo[ i+dy[a] ][ j+dx[a] ] == 0:

				dfs(i+dy[a], j+dx[a])

	return



h, w = map(int, input().split())

c = [["" for _ in range(w)] for _ in range(h)]

memo = [[0 for _ in range(w)] for _ in range(h)]



for i in range(h):

	c[i] = list(input())

	for j in range(w):

		if c[i][j] == "s":

			y, x = i, j



flg = False

dx = [1, 0, -1, 0]

dy = [0, 1, 0, -1]



dfs(y, x)



print("No")

どうしても分からなかったので、以下のQiitaで紹介されている解答コードをコピペして提出してみたのですが、寧ろTLEとなるケースが増える結果となりました。

蟻本をPythonで (初級編) - Qiita

なんとなくですが、どうもタッチの差でTLEになっているっぽいので少し工夫を考えてみます

理由が分かりました。
今までずっと PyPy3 (2.4.0) でコードを提出していたのですが、Python3 (3.4.3) でコードを提出するとACしました。

詳細は以下のQiitaが詳しいのですが、再帰処理と文字列操作に関してはPyPyだと結構遅くなってしまうようです。

qiita.com

ARC 031 B 埋め立て

B - 埋め立て

これは、Python的な意味でかなり難問でした。
まず、配列を普通に = 演算子でコピーしようとすると、関数内で参照渡しになってしまうということ。

つまり、以下のコードの field[i][j] だけ書き換えたいのに A[i][j] も連動して書き換わってしまいます。
なので、copy ライブラリを使っています。

ここにも落とし穴があり、二次元配列の時は copy.copy() ではなく、copy.deepcopy() を使わないと結局参照渡しになってしまいます。
10 x 10 の地図なので、4重ループを使うというのもここまで競プロに慣れてくると抵抗があるのも結構しんどいですね。

import copy



def dfs(x, y):

    field[x][y] = "x"



    for i in range(4):

        nx = x + dx[i]

        ny = y + dy[i]



        if 0 <= nx < 10 and 0 <= ny < 10 and field[nx][ny] == "o":

            dfs(nx, ny)



A = [list(input()) for i in range(10)]



dx = [1, 0, -1, 0]

dy = [0, 1, 0, -1]



for p in range(10):

    for q in range(10):

        if A[p][q] == "x":

            field = copy.deepcopy(A)

            field[p][q] = "o"

            count = 0

            for i in range(10):

                for j in range(10):

                    if field[i][j] == "o":

                        dfs(i, j)

                        count += 1

            if count == 1:

                print("YES")

                exit(0)



print("NO")

ARC 037 B バウムテスト

B - バウムテスト 

気付いてしまえばなんてことない問題ですが、解けませんでした。
関数dfsに前にいた値を持たせるところが思いつかなかったですね...!

def dfs(now, prev):

    global flag



    for next in g[now]:

        if next != prev:

            if memo[next]:

                flag = False

            else:

                memo[next] = True

                dfs(next, now)



n, m = map(int, input().split())

g = [[] for _ in range(n)]



for i in range(m):

    u, v = map(int, input().split())

    u -= 1

    v -= 1

    g[u].append(v)

    g[v].append(u)



memo = [False for i in range(n)]



ans = 0

for i in range(n):

    if not memo[i]:

        flag = True

        dfs(i, -1)

        if flag:

            ans += 1



print(ans)

AOJ 1160 島はいくつある?

How Many Islands?

これはほぼほぼ、蟻本と同じですね。以下になりました。

import sys

sys.setrecursionlimit(10**8)



ans = []



def dfs(x, y):

    c[x][y] = 0



    for dx in range(-1, 2):

        for dy in range(-1, 2):

            nx = x + dx

            ny = y + dy



            if 0 <= nx < h and 0 <= ny < w and c[nx][ny] == 1:

                dfs(nx, ny)



while True:

    w, h = map(int, input().split())

    if w == 0 and h ==0:

        break



    count = 0



    c = [[] for _ in range(h)]

    for i in range(h):

        if w == 1:

            c[i] = [int(input())]

        else:

            c[i] = list(map(int, input().split()))



    for x in range(h):

        for y in range(w):

            if c[x][y] == 1:

                dfs(x, y)

                count += 1



    ans.append(count)



for i in ans:

    print(i)

迷路の最短路(蟻本 P.37)

これは、Pythonのqueueライブラリを活用します。
そのまま本の通りに書けますね。

import queue



INF = 10 ** 8



n, m = map(int, input().split())

maze = ["" for _ in range(n)]



for i in range(n):

    maze[i] = list(input())



for i in range(n):

    for j in range(m):

        if maze[i][j] == "S":

            sx, sy = i, j

        elif maze[i][j] == "G":

            gx, gy = i, j



dx = [1, 0, -1, 0]

dy = [0, 1, 0, -1]



def bfs():

    que = queue.Queue()

    d = [[INF for _ in range(m)] for _ in range(n)]

    que.put((sx, sy))

    d[sx][sy] = 0



    while que.empty() == False:

        p = que.get()

        if p[0] == gx and p[1] == gy:

            break



        for i in range(4):

            nx = p[0] + dx[i]

            ny = p[1] + dy[i]



            if 0 <= nx < n and 0 <= ny < m:

                if maze[nx][ny] != "#" and d[nx][ny] == INF:

                    que.put((nx, ny))

                    d[nx][ny] = d[p[0]][p[1]] + 1



    return d[gx][gy]



res = bfs()

print(res)

ABC 007 C 幅優先探索

C - 幅優先探索

これはほぼ、蟻本と同じですね。
個人的に詰まったところは、 q.qut( (sy, sx) ) の内側の括弧を記載忘れてしまってコードが上手く動かず、TypeError: 'int' object is not subscriptable が出てしまうシーンでした。

import queue

INF = 10**8



def bfs():

    q = queue.Queue()

    q.put((sy, sx))

    d[sy][sx] = 0



    while q.empty() == False:

        p = q.get()

        if p == (gy, gx):

            break



        for i in range(4):

            ny = p[0] + dy[i]

            nx = p[1] + dx[i]



            if 0 <= ny < r and 0 <= nx < w and c[ny][nx] == "." and d[ny][nx] == INF:

                q.put((ny, nx))

                d[ny][nx] = d[p[0]][p[1]] + 1



r, w = map(int, input().split())

sy, sx = map(int, input().split())

gy, gx = map(int, input().split())

sy -= 1

sx -= 1

gy -= 1

gx -= 1



c = ["" for i in range(r)]

for i in range(r):

    c[i] = list(input())



dx = [1, 0, -1, 0]

dy = [0, 1, 0, -1]

d = [[INF for _ in range(w)] for _ in range(r)]



bfs()

print(d[gy][gx])

JOI 2010 予選 E チーズ

E - チーズ (Cheese)

これは、上記の問題のスタート地点とゴール地点を入れ替えて繰り返す形の問題ですね。
素直にこんな感じで解けました。
(計算量を考えることがだんだん難しくなってきました。)

import queue

INF = 10**8



def bfs(sx, sy, gx, gy):

    d = [[INF]*w for _ in range(h)]

    hp = 1

    q = queue.Queue()

    q.put((sx, sy))

    d[sx][sy] = 0



    while q.empty() == False:

        now = q.get()

        if now[0] == gx and now[1] == gy:

            return d[gx][gy]



        for i in range(4):

            nx = now[0] + dx[i]

            ny = now[1] + dy[i]



            if 0 <= nx < h and 0 <= ny < w and atlas[nx][ny] != "X" and d[nx][ny] == INF:

                q.put((nx, ny))

                d[nx][ny] = d[now[0]][now[1]] + 1



h, w, n = map(int, input().split())

atlas = [[] for i in range(h)]

cheese = [[] for i in range(n + 1)]



for i in range(h):

    atlas[i] = list(input())



for i in range(h):

    for j in range(w):

        if atlas[i][j] == "S":

            cheese[0] = (i, j)

        for k in range(1, n + 1):

            if atlas[i][j] == str(k):

                cheese[k] = (i, j)



dx = [1, 0, -1, 0]

dy = [0, 1, 0, -1]

count = 0



for i in range(n):

    count += bfs(cheese[i][0], cheese[i][1], cheese[i + 1][0], cheese[i + 1][1])



print(count)

ABC 088 D Grid Repainting

D - Grid Repainting

これは、スタート地点が (0, 0) で、ゴール地点が (h-1, w-1) として上記までの迷路問題を適用すると、以下のようになります。

  • ゴール地点までの最小移動回数が INF のままだったら、ゴールに辿り着けない
  • ゴール地点までの最小移動回数 + 1の数のマスだけ白くしておけば、後は全部黒く塗りつぶして良い

上記のアイディア + 事前に黒く塗り潰されている点の数をさらに引き算すると以下のようなコードになりました。

import queue

INF = 10**8



def bfs():

    q = queue.Queue()

    q.put((sx, sy))

    d[sx][sy] = 0



    while q.empty() == False:

        now = q.get()

        if now[0] == gx and now[1] == gy:

            break



        for i in range(4):

            nx = now[0] + dx[i]

            ny = now[1] + dy[i]



            if 0 <= nx < h and 0 <= ny < w and s[nx][ny] == "." and d[nx][ny] == INF:

                q.put((nx, ny))

                d[nx][ny] = d[now[0]][now[1]] + 1



h, w = map(int, input().split())

count = 0

s = [[] for i in range(h)]

for i in range(h):

    s[i] = list(input())

    for j in range(w):

        if s[i][j] == "#":

            count += 1



sx, sy = 0, 0

gx, gy = h - 1, w - 1

dx = [1, 0, -1, 0]

dy = [0, 1, 0, -1]

d = [[INF]*w for i in range(h)]



bfs()

if d[gx][gy] == INF:

    print(-1)

else:

    ans = h * w - (d[gx][gy] + 1 + count)

    print(ans)

AGC 033 A - Darker and Darker

A - Darker and Darker

これまでの流れを踏まえて素直に解くと以下のコードになるのかなと思います。
が、TLEしてしまいました。

import queue



def penetrate(x, y, count):

    global max_idx

    for i in range(4):

        nx = x + dx[i]

        ny = y + dy[i]



        if 0 <= nx < h and 0 <= ny < w and A[nx][ny] == "." and d[nx][ny] == -1:

            d[nx][ny] = count + 1

            q.put((nx, ny, d[nx][ny]))

            max_idx = max(max_idx, d[nx][ny])





h, w = map(int, input().split())

d = [[-1] * w for _ in range(h)]



A = [[] for i in range(h)]

q = queue.Queue()



for i in range(h):

    A[i] = list(input())

    for j in range(w):

        if A[i][j] == "#":

            d[i][j] = 0

            q.put((i, j, d[i][j]))



dx = [1, 0, -1, 0]

dy = [0, 1, 0, -1]

max_idx = 0



while q.empty() == False:

    idx = q.get()

    penetrate(idx[0], idx[1], idx[2])



print(max_idx)

ここまで、まともに計算量を考えてこなかったので改めて考えてみようと思います。
まず、for 文の二重ループに関しては O(H * W) なので 1.0e6 です。
うーん、結構この時点でギリギリですね。

次に、bfsを呼び出しているwhile文ですがこれは全マスについて計算しているので O(H * W) は確定です。
bfsの中で for i in range(4) としているので現時点で 4.0e6 です。

結構この時点でギリギリ感ありますが、逆に言えば queue 処理の計算量によってはなんとかなりそうなことも分かります。

ということで、調べてみたところ、queue.Queue() よりも collections.deque() を使った方が早いとのこと。
ただし、queue.Queue() は複数のスレッドを扱う場合でも安全に使うことができるのが利点のようです。

dequeに関して両端のアクセスはO(1)ということでした。
※queueに関してもたぶんそうですが。

O(1)なら、処理が早いdequeで勝てるのでは?
ということで collections.deque() を使って書き直してみました。

結果、PyPy3 で提出すれば無事ACとなりました。(Python3ではTLEする)
今後は queue じゃなくて deque を使うと誓いました。

from collections import deque



def penetrate(x, y, count):

    global max_idx

    for i in range(4):

        nx = x + dx[i]

        ny = y + dy[i]



        if 0 <= nx < h and 0 <= ny < w and A[nx][ny] == "." and d[nx][ny] == -1:

            d[nx][ny] = count + 1

            q.append((nx, ny, d[nx][ny]))

            max_idx = max(max_idx, d[nx][ny])





h, w = map(int, input().split())

d = [[-1] * w for _ in range(h)]



A = [[] for i in range(h)]

q = deque()



for i in range(h):

    A[i] = list(input())

    for j in range(w):

        if A[i][j] == "#":

            d[i][j] = 0

            q.append((i, j, d[i][j]))



dx = [1, 0, -1, 0]

dy = [0, 1, 0, -1]

max_idx = 0



while len(q) > 0:

    idx = q.popleft()

    penetrate(idx[0], idx[1], idx[2])



print(max_idx)

ARC 005 C 器物損壊!高橋君

C - 器物損壊!高橋君

タイトルめっちゃ面白いじゃん。
これまでの迷路について、迷路の塀を二回まで破壊できるオプション付きの問題です。
普通のbfsだと、計算量は O(H * W) なので今回の問題だと 2.5e3程度 になりそうですね。

ここに更に塀を2回壊したいです。
壊す場合と壊さない場合で考えると、1回だけでも 2 ^ (H * W) でもうダメです。
他の方法を考えないといけない。。

ここで、今回求めたいのは「高橋君が魚屋に辿り着くことができるか」に着目してみます。
一番最初の蟻本でたどり着くための最小移動回数を入れていた d というリストに、高橋くんの残り体力を格納していくことにします。

これなら、G が初期値かそうじゃないかでなんとか解けそう...!ということでやってみました。
PyPy3 だとMLEとなってしまいますが、Python3 での提出なら無事ACできました。嬉しい。

from collections import deque



def bfs(hp):

    q = deque()

    q.append((sx, sy, hp))

    d[sx][sy] = initial_hp



    while len(q) > 0:

        now = q.popleft()

        if now[0] == gx and now[1] == gy:

            break



        for i in range(4):

            nx = now[0] + dx[i]

            ny = now[1] + dy[i]



            if 0 <= nx < h and 0 <= ny < w:

                if d[nx][ny] == -1:

                    if c[nx][ny] != "#":

                        q.append((nx, ny, now[2]))

                        d[nx][ny] = now[2]

                    elif c[nx][ny] == "#":

                        if now[2] > 0:

                            q.append((nx, ny, now[2] - 1))

                            d[nx][ny] = now[2] - 1

                else:

                    if now[2] > d[nx][ny] and c[nx][ny] != "#":

                        d[nx][ny] = now[2]

                        q.append((nx, ny, now[2]))

                    elif now[2] - 1 > d[nx][ny] and c[nx][ny] == "#" and now[2] > 0:

                        d[nx][ny] = now[2] - 1

                        q.append((nx, ny, now[2] - 1))





h, w = map(int, input().split())

c = [[] for _ in range(h)]



for i in range(h):

    c[i] = list(input())

    for j in range(w):

        if c[i][j] == "s":

            sx, sy = i, j

        elif c[i][j] == "g":

            gx, gy = i, j



dx = [1, 0, -1, 0]

dy = [0, 1, 0, -1]

d = [[-1]*w for _ in range(h)]

initial_hp = 2



bfs(initial_hp)

if d[gx][gy] > -1:

    print("YES")

else:

    print("NO")

AOJ 0503 Cup

Cup

今までと比較すると毛色の違う問題です。
紙に書いてみると、各トレイをスタックと見做してカップ移動を考えられることが分かります。

そこで、トレイ3つの状態をキューとして渡すことでかなり汚いながらもコードは書けました。
が、案の定TLEしてしまったので解答を探して自分なり咀嚼して書いてみました。

アルゴリズムは自分が実装したものと同じでしたが、かなり分かりやすいコードで描かれており、かつ無駄な計算が少なかったです。
単純に自分には実装力が足りなかったのだなぁと凄く勉強になります。

特に参考になったのは以下の4点です。

  1. def main() を使うとそれだけで速くなる
  2. _, *list = [] って書き方マジで凄い(自分はこのためだけにfor文書いてました)
  3. 直前の操作を1桁の数字で表現してtに格納している(自分は直前のトレイの配列を入れていたので激重でした)
  4. q.pop()の結果は、カンマ区切りで複数の変数に格納できる
from collections import deque



def main():

    while True:

        n, m = [int(i) for i in input().split()]

        if n == 0 and m == 0:

            return

        _, *a = [int(i) for i in input().split()]

        _, *b = [int(i) for i in input().split()]

        _, *c = [int(i) for i in input().split()]

        a.insert(0, 0)

        b.insert(0, 0)

        c.insert(0, 0)



        q = deque()

        q.appendleft([a, b, c, 0, -1])

        tmp = [i for i in range(0, n+1)]



        while q:

            # d はカウンタ, tは操作の種類を示す

            a, b, c, d, t = q.pop()

            """

            print(a)

            print(b)

            print(c)

            print('=======')

            """

            # 終了

            if d > m:

                print(-1)

                break

            if a == tmp or c == tmp:

                print(d)

                break

            # a から b。

            # 数字の大小比較と直前の操作が「b から a」ではないを確認。

            if a[-1] > b[-1] and t != 1:

                q.appendleft([a[:-1], b+[a[-1]], c[:], d+1, 0])

            # b から a

            if b[-1] > a[-1] and t != 0:

                q.appendleft([a+[b[-1]], b[:-1], c[:], d+1, 1])

            # b から c

            if b[-1] > c[-1] and t != 3:

                q.appendleft([a[:], b[:-1], c+[b[-1]], d+1, 2])

            # c から b

            if c[-1] > b[-1] and t != 2:

                q.appendleft([a[:], b+[c[-1]], c[:-1], d+1, 3])



if __name__ == '__main__':

    main()

特殊な状態の列挙(蟻本 P39~40)

ABC 054 C One-stroke Path

C - One-stroke Path

分からなかったので、解答を読んでしまいました。
答えとなりうるルートを予め全部生成して、正しいかどうかを判断するんですね...!
なるほど。。。

import itertools



n, m = map(int, input().split())



path = [[False] * n for i in range(n)]

for i in range(m):

    a, b = map(int, input().split())

    a -= 1

    b -= 1

    path[a][b] = True

    path[b][a] = True



ans = 0



# 頂点を並び替える順列を生成してループ

for i in itertools.permutations(range(n), n):

    # 頂点1が始点

    if i[0] == 0:

        # 生成した順列の中をさらにループ

        for j in range(n):

            # n - 1 まで続いたら条件を満たすパスが存在する

            if j == n - 1:

                ans += 1

                break

            # i[j] から i[j + 1] に行くパスがなければ終了

            if not path[i[j]][i[j + 1]]:

                break



print(ans)

JOI 2009 予選 D カード並べ

Lining up the cards

非常に分かりやすい問題だった気がしています!
素直に解けますね。

import itertools



ans = []



while True:

	n = int(input())

	k = int(input())

	if n == 0 and k == 0:

		break



	cards = []

	ans_set = []



	for i in range(n):

		cards.append(input())



	for i in itertools.permutations(cards, k):

		ans_set.append(int("".join(i)))

	

	ans.append(len(set(ans_set)))



for i in ans:

	print(i)

yukicoder No.133 カードゲーム

No.133 カードゲーム - yukicoder

これも素直に解くことができますね。

import itertools



n = int(input())

A = list(map(int, input().split()))

B = list(map(int, input().split()))

a_rate_numerator = 0

count = 0



for a in itertools.permutations(A, n):

	for b in itertools.permutations(B, n):

		count += 1

		a_score = 0

		b_score = 0

		for i in range(n):

			if a[i] > b[i]:

				a_score += 1

			elif b[i] > a[i]:

				b_score += 1

		if a_score > b_score:

			a_rate_numerator += 1



print(a_rate_numerator / count)

まとめ

今回、かなり長くなりましたが一旦まとめます。
(特殊な状態の列挙、に関しては別記事にします。)

  1. 部分和問題のような、「〇〇する or しない」の二択で条件分岐するものはdfsで解けることが多い
  2. dfsを再帰関数で表現する場合、PyPyを使うと死ぬ
  3. 直前の処理を記憶させた変数を引数にしたりキューに渡すと良いことがある。
  4. 3の、直前の処理の記憶は判別できればなんでもいい。(1桁の数字でも良い)
  5. nの取り得る範囲が小さい時、itertoolsを使うと便利なことがある
  6. リストの取り扱い、下手にやると参照渡しになってしまうので注意。どうしてもの場合は、copy関数を使うと良い
  7. queue.Queue()よりもcollections.deque()を使うべき
  8. 高速化のテクニックは細かく色々ある

ということで、本日は以上です。
それではまた〜。

【Python】二分探索を行うbisectの関数の挙動について

この記事では、Pythonで二分探索を行うbisectライブラリについて記載します。

そもそも二分探索とは

二分探索については以下の記事がわかりやすいです。

codezine.jp

ざっくりまとめると、以下のような感じです。

  1. 数字がn個入ったリストRの中から、xを見つけ出したい
  2. まともに検索すると、Rを最初から最後まで全部確認しないといけない
  3. ここで、Rが昇順に整列しているとする
  4. Rの真ん中の数字がxより小さいならRの前半、xより大きいならRの後半にxが存在する(探す範囲を一気に1/2にできる)
  5. 半分になったRについても真ん中をxと比較してどんどん半分に絞りこんでいく

こんな感じです。
普通にやると、O(n) の計算量がかかる検索を O(log n) でできる感じです。

この二分探索について、Pythonではbisectというライブラリをインポートすることで自分で実装せずに使えるようになります。
が、ちょっとトリッキーな挙動で、公式リファレンスを読んでも分かりにくいのでまとめてみました。

docs.python.org

bisect関数の挙動

bisect関数は bisect(R: リスト, x: 検索したい数) みたいな感じで使います。
戻り値が結構独特で、「R に x を辞書順で挿入すべき位置を戻す」という仕様になっています。
つまり、Rにxが含まれていようがいまいが何かしら値が戻ってくる関数になっています。

詳細は以下のサンプルコードを読んでみてもらうと早いのですが、x を挿入できる位置の戻し方よって2つの関数があります。

  • 左側の位置を戻す bisect_left 関数
  • 右側の位置を戻す bisect_right 関数 ※名前を省略して bisect 関数としても良い

とにかく、以下のサンプルコードをご覧ください。

import bisect
a=[1,2,2,2,3,6]

print(bisect.bisect_left(a, 1))
# 0が戻り値

print(bisect.bisect_left(a, 4))
# 5が戻り値

print(bisect.bisect_left(a, 7))
# 6が戻り値

print(bisect.bisect_left(a, 2))
# 1が戻り値

print(bisect.bisect_right(a, 2))
# 4が戻り値

print(bisect.bisect_right(a, 4))
# 5が戻り値

print(bisect.bisect_right(a, 7))
# 6が戻り値

戻り値については、リストのスライス記法と同じで a=[1,2,2,2,3,6] の中で【2を挿入できる場所】にある , が左から何番目のものかを戻す。みたいな覚え方をしています。
leftだったら1つ目の , なので 1 を、rightだったら4つ目の , なので 4を戻す感じです。

別の捉え方として、以下の考え方も使う機会がそこそこありそうです。

  • bisect_left(R, x) : リストRに含まれるx未満の数値の個数
  • bisect_right(R, x): リストRに含まれるx以下の数値の個数

こんな感じで結構独特なので、例えばリストにxが含まれているかどうかを知りたい時には少し工夫をする必要があります。

例えば、以下のような形で使うことが多いのかなぁと思います。

import bisect
a=[1,2,2,2,3,6]

# xがaに含まれているかを確認したい
x = 2
i = bisect.bisect(a, x)
if a[i - 1] == x:
  print("Yes")
else:
  print("No")

x = 5
i = bisect.bisect(a, x)
if a[i - 1] == x:
  print("Yes")
else:
  print("No")

# x以上の数字がaに含まれているかどうかを確認したい
x = 7
i = bisect.bisect(a, x)
if i == len(a):
  print("Yes")
else:
  print("No")

x = 2
i = bisect.bisect(a, x)
if i == len(a):
  print("Yes")
else:
  print("No")

ちなみに、bisect.bisect_right と bisect.bisect は以下のように常に同じ挙動をします。
可読性を気にせず、自分だけで使う分にはbisect_right を使う必要はないですね。

import bisect
a=[1,2,2,2,3]

print(bisect.bisect_right(a, 2))
# 4が戻り値

print(bisect.bisect_right(a, 4))
# 5が戻り値

print(bisect.bisect(a, 2))
# 4が戻り値

print(bisect.bisect(a, 4))
# 5が戻り値

どんな時に二分探索を使う?

さて、ここまで紹介してきた二分探索の bisect 関数ですが、Rが整列されていない場合は一度整列する必要があります。
整列の計算量は一般に O(n) なので、一見わざわざ整列してまで使う意味がないように感じます。

が、例えば検索したいxをn通り試したい場合だとO(n^2) かかるところを、事前に整列しておくことで O(n * log n) で計算できます。
※nが大きい場合、事前の整列にかかる計算量 O(n) を考えても早く計算できる

具体的には、以下の記事の最後の問題で使ったりします。

www.yasu-p.com

ということで、本日は二分探索についての記事でした。
それでは〜。