214
202

More than 3 years have passed since last update.

蟻本をPythonで (初級編)

Last updated at Posted at 2019-05-26

はじめに

蟻本(プログラミングコンテストチャレンジブック)は C++ で書かれているので、それを Python で書きました。詳しい解説は本に書いてあるため省略します。蟻本片手にどうぞ。また、AtCoder 版!蟻本 (初級編)に載っている類題についても扱います。類題については先にリンク先で問題を見てから読むことを推奨します。
現在 2-3 までを載せています(少しずつ更新していきます)。

競技プログラミングをやったことがない方や、 Python の書き方がわからない方は以下の記事も参考にしてください。

2 初級編

2-1 すべての基本 "全探索"

スタックとキュー

Python では list 型の pop メソッドを使うことで以下のようにスタックとキューを扱うことができます。


# スタック
s = []
s.append(1)  # [1]
s.append(2)  # [1, 2]
s.append(3)  # [1, 2, 3]
s.pop()      # 一番上から取り除く [1, 2, 3] -> [1, 2]
s.pop()      # [1, 2] -> [1]
s.pop()      # [1] -> []


# キュー
q = []
q.append(1)  # [1]
q.append(2)  # [1, 2]
q.append(3)  # [1, 2, 3]
q.pop(0)     # 一番下から取り除く [1, 2, 3] -> [2, 3]
q.pop(0)     # [2, 3] -> [3]
q.pop(0)     # [3] -> []

しかし、list 型だと 0 番目の要素を取り除く際に O(n)かかってしまいます。そこで、この操作を O(1) で行うことができる collections.deque を利用します。

スタック


from collections import deque

s = deque([])
s.append(1)  # [1]
s.append(2)  # [1, 2]
s.append(3)  # [1, 2, 3]
s.pop()      # 一番上から取り除く [1, 2, 3] -> [1, 2]
s.pop()      # [1, 2] -> [1]
s.pop()      # [1] -> []

キュー

スタックと同様に collections.deque で扱うことができます。スタックでは取り除くときに pop だったのが、キューだと popleft になっているだけです。


from collections import deque

q = deque([])
q.append(1)  # [1]
q.append(2)  # [1, 2]
q.append(3)  # [1, 2, 3]
q.popleft()  # 一番下から取り除く [1, 2, 3] -> [2, 3]
q.popleft()  # [2, 3] -> [3]
q.popleft()  # [3] -> []

例題 2-1-1 部分和問題

部分和問題


# 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


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

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

実はこの問題は再帰以外に bit で解くこともでき、以下のようになります。


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

# 0 から (2^n)-1 までループ
for bit in range(1 << n):
    sum = 0

    for i in range(n):
        # bit に i 番目のフラグが立っているかどうか
        if bit & (1 << i):
            # フラグが立っているならば sum に加算する
            sum += a[i]

    if sum == k:
        print("Yes")
        exit()

print("No")

bit 全探索についてはビット演算 (bit 演算) の使い方を総特集! 〜 マスクビットから bit DP まで 〜を参考にしました。

ABC 045 C たくさんの数式

・再帰


def dfs(i, f):
    if i == n - 1:
        return sum(list(map(int, f.split("+"))))

    # 式 f の末尾に "+" を追加するかしないかして次の数字を追加
    return dfs(i + 1, f + s[i + 1]) + \
           dfs(i + 1, f + "+" + s[i + 1])


s = input()
n = len(s)

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

・bit

s = input()
n = len(s)

ans = 0

for bit in range(1 << (n - 1)):
    # 各場合で式 f を生成する
    f = s[0]

    for i in range(n - 1):
        if bit & (1 << i):
            # フラグが立っているならば "+" を式の末尾に追加する
            f += "+"
        f += s[i + 1]

    ans += sum(map(int, f.split("+")))

print(ans)

ABC 079 C Train Ticket

ABC 045 C たくさんの数式とほぼ同じです。

・再帰


def dfs(i, f, sum):
    if i == 3:
        if sum == 7:
            # 答えは1つ出力すれば良いので =7 になれば終了
            print(f + "=7")
            exit()

    else:
        # 式 f の末尾に符号と次の数字を追加し、その分 sum に加減する
        dfs(i + 1, f + "-" + s[i + 1], sum - int(s[i + 1]))
        dfs(i + 1, f + "+" + s[i + 1], sum + int(s[i + 1]))


s = input()

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

・bit


s = input()

for bit in range(1 << 3):
    ans = int(s[0])
    f = s[0]

    for i in range(3):
        # フラグが立っていれば "+" 、なければ "-"
        if bit & (1 << i):
            ans += int(s[i + 1])
            f += "+"
        else:
            ans -= int(s[i + 1])
            f += "-"
        f += s[i + 1]

    if ans == 7:
        print(f + "=7")
        exit()

ABC 104 C All Green

中途半端に解く配点は1種類以下で良いので、各配点について解くか解かないかを決め、その上で G 点に満たなければまだ解いてない配点のうち一番高い配点のものを足りない分解けば良いことになります。

・再帰


# まだ解いてない配点を nokori として持つ
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)

・bit


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

ans = float("inf")

for bit in range(1 << d):
    count = 0
    sum = 0
    nokori = set(range(1, d + 1))

    for i in range(d):
        if bit & (1 << i):
            sum += pc[i][0] * (i + 1) * 100 + pc[i][1]
            count += pc[i][0]
            nokori.discard(i + 1)

    # 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)

print(ans)

ARC 029 A 高橋君とお肉

・再帰


def dfs(i, a, b):
    global ans
    if i == n:
        ans = min(ans, max(a, b))

    else:
        # a の肉焼き器に乗せるか、 b の肉焼き器に乗せるか
        dfs(i + 1, a + t[i], b)
        dfs(i + 1, a, b + t[i])


n = int(input())
t = [int(input()) for i in range(n)]

ans = float("inf")

dfs(0, 0, 0)
print(ans)

・bit


n = int(input())
t = [int(input()) for i in range(n)]

ans = float("inf")

for bit in range(1 << n):
    a = 0
    b = 0

    for i in range(n):
        # フラグが立っていれば a に、なければ b に乗せる
        if bit & (1 << i):
            a += t[i]
        else:
            b += t[i]

    ans = min(ans, max(a, b))

print(ans)

ABC 002 D 派閥

・再帰


import itertools


# 形成した派閥を group として持つ
def dfs(i, group):
    global ans
    if i == n:
        # group 内の全員が知り合い同士か
        flag = True

        # group 内から2人選ぶ組み合わせのループ
        for i in itertools.combinations(group, 2):
            # 1人でも知り合いでなければ終了
            if friend[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])


n, m = map(int, input().split())
friend = [[0] * n for i in range(n)]
for i in range(m):
    x, y = map(int, input().split())
    x -= 1
    y -= 1
    friend[x][y] = 1
    friend[y][x] = 1

ans = 0

dfs(0, [])
print(ans)

・bit


import itertools

n, m = map(int, input().split())
friend = [[0] * n for i in range(n)]
for i in range(m):
    x, y = map(int, input().split())
    x -= 1
    y -= 1
    friend[x][y] = 1
    friend[y][x] = 1

ans = 0

for bit in range(1 << n):
    group = []

    for i in range(n):
        if bit & (1 << i):
            group.append(i)

    # group 内の全員が知り合い同士か
    flag = True

    # group 内から2人選ぶ組み合わせのループ
    for i in itertools.combinations(group, 2):
        # 1人でも知り合いでなければ終了
        if friend[i[0]][i[1]] == 0:
            flag = False
            break

    if flag:
        ans = max(ans, len(group))

print(ans)

例題 2-1-1 Lake Counting (POJ No.2386)

Lake Counting (POJ No.2386)


# 現在位置 (x, y)
def dfs(x, y):
    # 今いるところを . に置き換える
    field[x][y] = "."

    # 移動する8方向をループ
    for dx in range(-1, 2):
        for dy in range(-1, 2):
            # x, y 方向それぞれに dx, dy 移動した場所を (nx, ny) とする
            nx = x + dx
            ny = y + dy

            # nx と ny が庭の範囲内かどうかと field[nx][ny] が水溜りかどうかを判定
            if 0 <= nx and nx < n and 0 <= ny and ny < m and field[nx][ny] == "W":
                dfs(nx, ny)


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

res = 0
for i in range(n):
    for j in range(m):
        if field[i][j] == "W":
            # Wが残っているならそこから dfs をはじめる
            dfs(i, j)
            res += 1
print(res)

ATC 001 A 深さ優先探索


# pythonだとデフォルトの再帰処理回数の上限に引っかかってしまうので、それを変更
import sys
sys.setrecursionlimit(1000000)


def dfs(x, y):
    d[x][y] = 1

    # 移動4方向をループ
    for i in range(4):
        nx = x + dx[i]
        ny = y + dy[i]

        # nx と ny が街の範囲内か、行ったことがないか、塀ではないかを判定
        if 0 <= nx and nx < n and 0 <= ny and ny < m and d[nx][ny] == 0 and c[nx][ny] != "#":
                dfs(nx, ny)


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

# 到達したかどうか(0は未到達、1は到達済み)
d = [[0] * m for i in range(n)]

# 移動する4方向
dx = [1, 0, -1, 0]
dy = [0, 1, 0, -1]

# スタート地点から dfs を始める
for i in range(n):
    for j in range(m):
        if c[i][j] == "s":
            dfs(i, j)

# ゴール地点に到達したかどうか
for i in range(n):
    for j in range(m):
        if c[i][j] == "g" and d[i][j]:
            print("Yes")
            exit()
print("No")

問題のページにも解説が載っていますが、蟻本の書き方を採用しました。

ARC 031 B 埋め立て


def dfs(x, y):
    field[x][y] = "x"

    # 移動4方向をループ
    for i in range(4):
        nx = x + dx[i]
        ny = y + dy[i]

        # nx と ny が地図の範囲内か、field[nx][ny]が島か判定
        if 0 <= nx and nx < 10 and 0 <= ny and ny < 10 and field[nx][ny] == "o":
                dfs(nx, ny)


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

# 移動する4方向
dx = [1, 0, -1, 0]
dy = [0, 1, 0, -1]

# 埋め立てる1マスを全探索
for p in range(10):
    for q in range(10):
        if A[p][q] == "x":
            # 海だったら A を field にコピーしてそのマスを埋め立てる
            field = [k[:] for k in A]
            field[p][q] = "o"
            count = 0
            # 埋め立てた上で dfs で島の数をカウント
            for i in range(10):
                for j in range(10):
                    if field[i][j] == "o":
                        dfs(i, j)
                        count += 1
            # count (島の数)が1かどうか
            if count == 1:
                print("YES")
                exit()
print("NO")

ARC 037 B バウムテスト


def dfs(now, prev):
    global flag
    # 今いる頂点から行ける頂点を順に next に入れてループ
    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 = [[] * n for i 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)

例題 2-1-3 迷路の最短路

迷路の最短路


from collections import deque


# (sx, sy) から (gx, gy) への最短距離を求める
# 辿り着けないと INF
def bfs():
    # すべての点を INF で初期化
    d = [[float("inf")] * m for i in range(n)]

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

    for i in range(n):
        for j in range(m):
            # スタートとゴールの座標を探す
            if maze[i][j] == "S":
                sx = i
                sy = j
            if maze[i][j] == "G":
                gx = i
                gy = j

    # スタート地点をキューに入れ、その点の距離を0にする
    que = deque([])
    que.append((sx, sy))
    d[sx][sy] = 0

    # キューが空になるまでループ
    while que:
        # キューの先頭を取り出す
        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] == float("inf"):
                # 移動できるならキューに入れ、その点の距離を p からの距離 +1 で確定する
                que.append((nx, ny))
                d[nx][ny] = d[p[0]][p[1]] + 1

    return d[gx][gy]


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

print(bfs())

ABC 007 C 幅優先探索

蟻本の問題とほぼ全く同じですが、sx と sy 、gx と gy が逆なので注意が必要です。


from collections import deque


def bfs():
    d = [[float("inf")] * c for i in range(r)]

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

    que = deque([])
    que.append((sy, sx))
    d[sy][sx] = 0

    while que:
        p = que.popleft()
        if p[0] == gy and p[1] == gx:
            break
        for i in range(4):
            nx = p[0] + dx[i]
            ny = p[1] + dy[i]

            if 0 <= nx < r and 0 <= ny < c and maze[nx][ny] != "#" and d[nx][ny] == float("inf"):
                que.append((nx, ny))
                d[nx][ny] = d[p[0]][p[1]] + 1

    return d[gy][gx]


r, c = map(int, input().split())
sy, sx = map(int, input().split())
gy, gx = map(int, input().split())
sx -= 1
sy -= 1
gy -= 1
gx -= 1
maze = [list(input()) for i in range(r)]

print(bfs())

JOI 2010 予選 E チーズ

チーズを食べるたびにスタートとゴールが変わる問題です。


from collections import deque


def bfs(sx, sy, gx, gy):
    d = [[float("inf")] * w for i in range(h)]

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

    que = deque([])
    que.append((sx, sy))
    d[sx][sy] = 0

    while que:
        p = que.popleft()
        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 < h and 0 <= ny < w and maze[nx][ny] != "X" and d[nx][ny] == float("inf"):
                que.append((nx, ny))
                d[nx][ny] = d[p[0]][p[1]] + 1

    return d[gx][gy]


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

# 各チーズ工場の座標
g = [0] * n

for i in range(h):
    for j in range(w):
        if maze[i][j] == "S":
            sx = i
            sy = j
        # 書かれているのが数字なら g にその座標を書き込む
        if maze[i][j].isdecimal():
            g[int(maze[i][j]) - 1] = (i, j)

# 1回目は巣をスタート、硬さ1のチーズ工場をゴールにする
ans = bfs(sx, sy, g[0][0], g[0][1])
# それ以降は前回の工場をスタート、次の工場をゴールにする
for i in range(1, n):
    ans += bfs(g[i - 1][0], g[i - 1][1], g[i][0], g[i][1])
print(ans)

ABC 088 D Grid Repainting

最短経路のマス以外は塗ることができるため、白いマスの個数から最短経路のマス分とゴールのマスを引けば求められます。


from collections import deque


def bfs():
    d = [[float("inf")] * w for i in range(h)]

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

    que = deque([])
    que.append((sx, sy))
    d[sx][sy] = 0

    while que:
        p = que.popleft()
        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 < h and 0 <= ny < w and maze[nx][ny] != "#" and d[nx][ny] == float("inf"):
                que.append((nx, ny))
                d[nx][ny] = d[p[0]][p[1]] + 1

    return d[gx][gy]


h, w = map(int, input().split())
maze = [list(input()) for i in range(h)]
sx, sy = 0, 0
gx, gy = h - 1, w - 1

# 白いマスを数える
white = 0
for i in range(h):
    for j in range(w):
        if maze[i][j] == ".":
            white += 1

res = bfs()
if 0 < res < float("inf"):
    # 白いマスの数から最短経路でかかるコスト分を引く
    # ゴールを黒いマスに変えることはできないためその分も引く
    print(white - res - 1)
else:
    print(-1)

ARC 005 C 器物損壊!高橋君

壁を壊した回数も含めて考える必要があるため、今までの d[x][y] を d[x][y][t](t は壊した回数) として解きます。また、 Python だと BFS で解くと TLE になってしまうため、今回は DFS (deque の popleft を pop に変えただけです)を使いました。


from collections import deque


def dfs():
    d = [[[False] * 3 for j in range(w)] for i in range(h)]

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

    for i in range(h):
        for j in range(w):
            if maze[i][j] == "s":
                sx = i
                sy = j
            if maze[i][j] == "g":
                gx = i
                gy = j

    que = deque([])
    # 3つ目は壊した回数
    que.append((sx, sy, 0))
    # 座標と壊した回数ごとに行ける行けないを持つ
    d[sx][sy][0] = True

    while que:
        # ここが popleft ではなく pop になっているためスタックになり、これはDFSです
        p = que.pop()
        if p[0] == gx and p[1] == gy:
            break
        for i in range(4):
            nx = p[0] + dx[i]
            ny = p[1] + dy[i]
            t = p[2]

            if not (0 <= nx < h and 0 <= ny < w):
                continue
            # すでに2回壊してる状態で壁に当たると進めない
            if t == 2 and maze[nx][ny] == "#":
                continue
            if d[nx][ny][t]:
                continue
            # 壁なら壊して進む
            if maze[nx][ny] == "#":
                que.append((nx, ny, t + 1))
                d[nx][ny][t + 1] = True
            else:
                que.append((nx, ny, t))
                d[nx][ny][t] = True

    # ゴールの座標を壊した回数でループ
    for i in range(3):
        if d[gx][gy][i]:
            return True


h, w = map(int, input().split())
maze = [list(input()) for i in range(h)]

if dfs():
    print("YES")
else:
    print("NO")

例題 2-1-4 特殊な状態の列挙

特殊な状態の列挙

Python には順列や組み合わせを列挙することのできる itertools というライブラリがあります。


import itertools


# 順列
# permutations(list, n) で list から n 個選んで並べる
for i in itertools.permutations([0, 1, 2], 3):
    print(i)
# (0, 1, 2)
# (0, 2, 1)
# (1, 0, 2)
# (1, 2, 0)
# (2, 0, 1)
# (2, 1, 0)


# 組み合わせ
# combinations(list, n) で list から n 個選ぶ
for i in itertools.combinations([0, 1, 2], 2):
    print(i)
# (0, 1)
# (0, 2)
# (1, 2)

ABC 054 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 カード並べ


import itertools

n = int(input())
k = int(input())
card = [input() for i in range(n)]

# set を使うと重複せずに数えられる
number = set()
# card から k 個選んで並び替える
for i in itertools.permutations(card, k):
    # 並び替えたものを繋げて1つの文字列にする
    number.add("".join(i))

print(len(number))

2-2 猪突猛進! "貪欲法"

例題 2-2-1 硬貨の問題

硬貨の問題


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

# コインの金額
v = [1, 5, 10, 50, 100, 500]
ans = 0

for i in range(1, 7):
    t = min(a // v[-i], c[-i])  # コイン i を使う枚数
    a -= t * v[-i]
    ans += t

print(ans)

JOI 2007 予選 A おつり

「硬貨の問題」の各硬貨を十分にたくさん持っているバージョンです。


# おつりの金額
a = 1000 - int(input())

v = [1, 5, 10, 50, 100, 500]
ans = 0

for i in range(1, 7):
    t = a // v[-i]
    a -= t * v[-i]
    ans += t

print(ans)

例題 2-2-2 区間スケジューリング問題

区間スケジューリング問題


# 多次元リストの sort に使える
from operator import itemgetter

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

st = sorted([(s[i], t[i]) for i in range(n)], key=itemgetter(1))

ans = 0
# 最後に選んだ仕事の終了時間
last = 0

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

print(ans)

KUPC 2015 A 東京都


t = int(input())

for i in range(t):
    s = input()
    ans = 0
    j = 0

    while j < len(s) - 4:
        # 前から5文字ずつ見ていく
        if s[j:j + 5] == "tokyo" or s[j:j + 5] == "kyoto":
            ans += 1
            # 一致すれば5文字分進める
            j += 5
        else:
            j += 1

    print(ans)

ABC 103 D - Islands War


from operator import itemgetter

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

# 区間の終端でソート
ab = sorted([tuple(map(int, input().split())) for i in range(m)], key=itemgetter(1))
# 前回除いた橋
removed = -1
ans = 0

for a, b in ab:
    # a が removed より大きい = まだ取り除いてない
    if a > removed:
        removed = b - 1
        ans += 1

print(ans)

ABC 038 D プレゼント

解けなくて調べてたら蟻本の中級編に載ってるデータ構造が必要だということがわかったので一旦飛ばします。

例題 2-2-3 Best Cow Line (POJ No.3617)

Best Cow Line (POJ No.3617)


n = int(input())
s = input()

t = ""

a = 0
b = 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


sd = input()
t = input()

n = len(sd)
m = len(t)
s = []

# sd を後ろから見ていき、 t の入りそうな場所を探す
for i in range(n - m, -1, -1):
    t_kamo = sd[i:i + m]
    for j in range(m + 1):
        # 1文字ずつ順に入りうるか調べ、最後まで入るなら "?" を "a" に置き換えて出力
        if j == m:
            print((sd[:i] + t + sd[i + len(t):]).replace("?", "a"))
            exit()
        if t_kamo[j] == "?":
            continue
        elif t_kamo[j] != t[j]:
            break

print("UNRESTORABLE")

ABC 007 B 辞書式順序

辞書順最小は "a" なので、これを出力すれば良いです。


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

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


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 の中で元の位置と違う文字の数
        diff1 = diff + (c != s[i])
        # まだ使ってない文字の中で元の位置と違う文字の数
        diff2 = counts - (counter[c] > 0)

        # 両方を足して 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)

例題 2-2-4 Saruman's Army (POJ No.3069)

Saruman's Army (POJ No.3069)


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

x = sorted(x)

i = 0
ans = 0
while i < n:
    # s はカバーされていない一番左の点の位置
    s = x[i]
    i += 1
    # s から距離 r を超える点まで進む
    while i < n and x[i] <= s + r:
        i += 1

    # p は新しく印を付ける点の位置
    p = x[i - 1]
    # p から距離 r を超える点まで進む
    while i < n and x[i] <= p + r:
        i += 1

    ans += 1

print(ans)

ABC 083 C Multiple Gift

条件を満たすもののうち数列の長さが最大になるのは、 Ai+1 が Ai の2倍になるような数列です。


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

ans = 0

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

print(ans)

ARC 006 C 積み重ね


n = int(input())
w = [int(input()) for i in range(n)]

ans = 0
# 積み重ねた山の一番上のダンボール
top = []

for i in w:
    # 山を順番に見る
    for j in range(len(top) + 1):
        # 載せられる山がなければ新しい山を作る
        if j == len(top):
            ans += 1
            top.append(i)
        # 載せられる山があればそこに載せる
        if top[j] >= i:
            top[j] = i
            break

print(ans)

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


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

# 各たこ焼きが作られてから経過した時間
takoyaki = []
# 今いる客の数
kyaku = 0

# 100秒目までループ
for i in range(101):
    # 客を待たせてはいけない
    if kyaku:
        print("no")
        exit()

    for j in range(n):
        # i 秒目にたこ焼きが作られていれば takoyaki に追加
        if a[j] == i:
            takoyaki.append(0)
    for k in range(m):
        # i 秒目に客が来ていれば kyaku に加算
        if b[k] == i:
            kyaku += 1

    # 客とたこ焼きがある限り売る
    while takoyaki and kyaku:
        takoyaki.pop(0)
        kyaku -= 1

    x = 0
    while x < len(takoyaki):
        # 各たこ焼きに 1 秒加算する
        takoyaki[x] += 1
        # t 秒を超えたら捨てる
        if takoyaki[x] > t:
            takoyaki.pop(x)
        else:
            x += 1

if kyaku:
    print("no")
else:
    print("yes")

例題 2-2-5 Fence Repair (POJ No.3253)

Fence Repair (POJ No.3253)


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

ans = 0

# 板が1本になるまで適用
while n > 1:
    # 一番短い板 mii1 、次に短い板 mii2 を求める
    mii1 = 0
    mii2 = 1

    if L[mii1] > L[mii2]:
        mii1, mii2 = mii2, mii1

    for i in range(2, n):
        if L[i] < L[mii1]:
            mii2 = mii1
            mii1 = i
        elif L[i] < L[mii2]:
            mii2 = i

    # それらを併合
    t = L[mii1] + L[mii2]
    ans += t

    if mii1 == n - 1:
        mii1, mii2 = mii2, mii1
    L[mii1] = t
    L[mii2] = L[n - 1]
    n -= 1

print(ans)

2-3 値を覚えて再利用 "動的計画法"

例題 2-3-1 01ナップサック問題

01ナップサック問題

まず、以下のコードは普通に全探索したものです。


# i 番目以降の品物から重さの総和が j 以下となるように選ぶ
def rec(i, j):
    if i == n:
        # もう品物は残っていない
        res = 0
    elif j < w[i]:
        # この品物は入らない
        res = rec(i + 1, j)
    else:
        # 入れない場合と入れる場合を両方試す
        res = max(rec(i + 1, j), rec(i + 1, j - w[i]) + v[i])

    return res


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

print(rec(0, W))

これでは最悪 $O(2^n)$ かかってしまうため、これを改善していきます。

  • メモ化

全探索の際に同じ計算が複数回行われているため、それをメモして再利用するメモ化を行います。


def rec_memo(i, j):
    if dp[i][j]:
        # すでに調べたことがあるならばその結果を再利用
        return dp[i][j]
    if i == n:
        res = 0
    elif j < w[i]:
        res = rec_memo(i + 1, j)
    else:
        res = max(rec_memo(i + 1, j), rec_memo(i + 1, j - w[i]) + v[i])
    # 結果をテーブルに記憶する
    dp[i][j] = res
    return res


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

dp = [[0] * (W + 1) for i in range(n + 1)]  # メモ化テーブル

print(rec_memo(0, W))

こうすることで計算量は $O(nW)$ と大幅に高速化できます。

  • 漸化式

また、同じ計算量でも漸化式を用いたループで解く方法もあります。


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

dp = [[0] * (W + 1) for i in range(n + 1)]  # DP テーブル

for i in range(n - 1, -1, -1):
    for j in range(W + 1):
        if j < w[i]:
            dp[i][j] = dp[i + 1][j]
        else:
            dp[i][j] = max(dp[i + 1][j], dp[i + 1][j - w[i]] + v[i])

print(dp[0][W])
  • 0 から見る漸化式

先ほどは上からループを回していきましたが、これを 0 から見ていくこともできます。


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

dp = [[0] * (W + 1) for i in range(n + 1)]

for i in range(n):
    for j in range(W + 1):
        if j < w[i]:
            dp[i + 1][j] = dp[i][j]
        else:
            dp[i + 1][j] = max(dp[i][j], dp[i][j - w[i]] + v[i])

print(dp[n][W])
  • 今の場所から次の場所へ遷移

また、今までは「今の場所にどこから遷移してこられるか」を見てきましたが、「今の場所からどこへ遷移できるか」を考えることもできます。


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

dp = [[0] * (W + 1) for i in range(n + 1)]

for i in range(n):
    for j in range(W + 1):
        dp[i + 1][j] = max(dp[i + 1][j], dp[i][j])
        if j + w[i] <= W:
            dp[i + 1][j + w[i]] = max(dp[i + 1][j + w[i]], dp[i][j] + v[i])

print(dp[n][W])

TDPC A コンテスト


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

dp = [False for i in range(10001)]  # ありうる得点は 100 * 100 = 10000 点まで

dp[0] = True  # 何も解いていなければ 0 点

for i in range(n):
    for j in range(10001 - p[i], -1, -1):  # 今回加算した分が影響しないように上から見ていく
        if dp[j]:  # j 点がありうるなら j + p[i] 点もありうる
            dp[j + p[i]] = True

print(sum(dp))

TDPC D サイコロ


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

# サイコロの目を素因数分解すると、2, 3, 5しか現れない
count2 = 0
count3 = 0
count5 = 0
while d % 2 == 0:
    d //= 2
    count2 += 1
while d % 3 == 0:
    d //= 3
    count3 += 1
while d % 5 == 0:
    d //= 5
    count5 += 1

# 2, 3, 5以外があったらだめ
if d != 1:
    print(0)
    exit()

dp = [[[[0] * (count5 + 1) for i in range(count3 + 1)] for j in range(count2 + 1)] for k in range(n + 1)]

dp[0][0][0][0] = 1

# サイコロの目1~6を素因数分解
d2 = [0, 1, 0, 2, 0, 1]
d3 = [0, 0, 1, 0, 0, 1]
d5 = [0, 0, 0, 0, 1, 0]

# dp[i][c2][c3][c5] := サイコロを i 回振って、2^c2 * 3^c3 + 5^c5 になる確率
for i in range(n):
    for c2 in range(count2 + 1):
        for c3 in range(count3 + 1):
            for c5 in range(count5 + 1):
                for j in range(6):
                    nc2 = min(count2, c2 + d2[j])
                    nc3 = min(count3, c3 + d3[j])
                    nc5 = min(count5, c5 + d5[j])
                    dp[i + 1][nc2][nc3][nc5] += dp[i][c2][c3][c5] / 6

print(dp[n][count2][count3][count5])

ABC 015 D 高橋くんの苦悩

この問題ですが、計算量が $O(NKW)$ なので C++ などでは間に合う計算量ですが、残念ながら Python では少し厳しいので PyPy で提出しないと AC できません。より良いアルゴリズムを使えば Python で通すことも可能ですが、ここでは割愛します。(同じコードで PyPy で通るならそれで良くないですか?)


w = int(input())
n, k = map(int, input().split())
a = []
b = []
for i in range(n):
    a_, b_ = map(int, input().split())
    a.append(a_)
    b.append(b_)

dp = [[[0] * (w + 1) for i in range(k + 1)] for j in range(n + 1)]

# dp[i][j][l] := i 番目まで見て、j 枚使用し、幅が合計 l のときの最大値
for i in range(n):
    for j in range(k + 1):
        for l in range(w + 1):
            if j == k:  # この枚数以上貼れない
                dp[i + 1][j][l] = dp[i][j][l]
            else:
                if l < a[i]:  # 幅が足りない
                    dp[i + 1][j][l] = dp[i][j][l]
                else:
                    dp[i + 1][j][l] = max(dp[i][j][l], dp[i][j - 1][l - a[i]] + b[i])

# 必ず k 枚使う必要はない
ans = 0
for i in range(k + 1):
    ans = max(ans, dp[n][i][w])

print(ans)

JOI 2012 予選 D 暑い日々

d, n = map(int, input().split())
t = [int(input()) for i in range(d)]
kion = []
hadesa = []
for i in range(n):
    a, b, c = map(int, input().split())
    kion.append((a, b))
    hadesa.append(c)

dp = [[-1] * n for i in range(d + 1)]

# dp[i][j] := i 日目に服 j を着たときの派手さの差の絶対値の合計の最大値

# 初日に着る服
for i in range(n):
    if kion[i][0] <= t[0] <= kion[i][1]:
        dp[1][i] = 0

for i in range(d):
    for j in range(n):
        if kion[j][0] <= t[i] <= kion[j][1]:  # 着るのに適している気温
            for k in range(n):  # 前日着た服の中で一番差が大きいものを選ぶ
                if dp[i][k] != -1:  # 着れなかったものを除く
                    dp[i + 1][j] = max(dp[i + 1][j], dp[i][k] + abs(hadesa[k] - hadesa[j]))

print(max(dp[d]))

例題 2-3-2 最長共通部分列問題

最長共通部分列問題


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

dp = [[0] * (m + 1) for i in range(n + 1)]

for i in range(n):
    for j in range(m):
        if s[i] == t[j]:
            dp[i + 1][j + 1] = dp[i][j] + 1
        else:
            dp[i + 1][j + 1] = max(dp[i][j + 1], dp[i + 1][j])

print(dp[n][m])

Indeedなう C Optimal Recommendations

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

# 能力値は 0 ~ 100 の 101 通り
dp = [[[0] * 102 for i in range(102)] for j in range(102)]

# dp[i][j][k] := 技術力 a, 語学力 j, コミュニケーション力 k のときの最高年収

for i in range(n):
    a, b, c, w = map(int, input().split())
    dp[a][b][c] = max(dp[a][b][c], w)

for i in range(101):  # 能力値は 0 ~ 100
    for j in range(101):
        for k in range(101):
            # 能力が高ければそれより要求能力の低い求人は応募できる
            dp[i + 1][j][k] = max(dp[i + 1][j][k], dp[i][j][k])
            dp[i][j + 1][k] = max(dp[i][j + 1][k], dp[i][j][k])
            dp[i][j][k + 1] = max(dp[i][j][k + 1], dp[i][j][k])

for i in range(m):
    x, y, z = map(int, input().split())
    print(dp[x][y][z])

このまま Python で提出すると TLE してしまうため、全体を関数内に入れるか、PyPy で提出すると通ります。全体を関数内に入れる場合は、以下のようにすれば大丈夫です。

def main():
    # ここに本来のコードを入れる

main()

例題 2-3-3 個数制限なしナップサック問題

個数制限なしナップサック問題

まずは、今までと同じような DP のコードを書いてみます。


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

dp = [[0] * (W + 1) for i in range(n + 1)]

for i in range(n):
    for j in range(W + 1):
        k = 0
        while k * w[i] <= j:
            dp[i + 1][j] = max(dp[i + 1][j], dp[i][j - k * w[i]] + k * v[i])
            k += 1

print(dp[n][W])

これだと三重ループになってしまうため、計算量は $O(nW^2)$ となります。これでは不十分なので、漸化式の式変形を行うことで k に関するループをなくすことができ、計算量を $O(nW)$ とすることができます。


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

dp = [[0] * (W + 1) for i in range(n + 1)]

for i in range(n):
    for j in range(W + 1):
        if j < w[i]:
            dp[i + 1][j] = dp[i][j]
        else:
            dp[i + 1][j] = max(dp[i][j], dp[i + 1][j - w[i]] + v[i])

print(dp[n][W])

また、最初の01ナップサック問題やこの個数制限なしナップサック問題は、同じ配列を再利用して更新していくことで解くこともできます。

・01ナップサック問題


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

dp = [0] * (W + 1)

for i in range(n):
    for j in range(W, w[i] - 1, -1):
        dp[j] = max(dp[j], dp[j - w[i]] + v[i])

print(dp[W])

どこが変わったかというと、配列が二次元から一次元になっています。配列の中身がどう変わってるかを、例題の入力の場合で見てみます。

-入力-
4
2 3
1 2
3 4
2 2
5
・二次元の場合の遷移
i=0, 0<=j<5
[[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]

i=1, 0<=j<5
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]

i=2, 0<=j<5
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 0, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 5, 0, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 5, 6, 0], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 5, 6, 7], [0, 0, 0, 0, 0, 0]]

i=3, 0<=j<5
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 5, 6, 7], [0, 0, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 5, 6, 7], [0, 2, 0, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 5, 6, 7], [0, 2, 3, 0, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 5, 6, 7], [0, 2, 3, 5, 0, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 5, 6, 7], [0, 2, 3, 5, 6, 0]]
[[0, 0, 0, 0, 0, 0], [0, 0, 3, 3, 3, 3], [0, 2, 3, 5, 5, 5], [0, 2, 3, 5, 6, 7], [0, 2, 3, 5, 6, 7]]
・一次元の場合の遷移
i=0, j=5,4,3,2
[0, 0, 0, 0, 0, 3]
[0, 0, 0, 0, 3, 3]
[0, 0, 0, 3, 3, 3]
[0, 0, 3, 3, 3, 3]

i=1, j=5,4,3,2,1
[0, 0, 3, 3, 3, 5]
[0, 0, 3, 3, 5, 5]
[0, 0, 3, 5, 5, 5]
[0, 0, 3, 5, 5, 5]
[0, 2, 3, 5, 5, 5]

i=2, j=5,4,3
[0, 2, 3, 5, 5, 7]
[0, 2, 3, 5, 6, 7]
[0, 2, 3, 5, 6, 7]

i=3, j=5,4,3,2
[0, 2, 3, 5, 6, 7]
[0, 2, 3, 5, 6, 7]
[0, 2, 3, 5, 6, 7]
[0, 2, 3, 5, 6, 7]

二次元の場合、更新のために使用してる値は一つ前の i だけです。つまり、前回更新し終えた配列にそのまま今回のものを上書きしていくことで、i のための配列を持つ必要がなくなります。そのため、一次元の遷移は二次元の遷移をぎゅっと縮めた感じになっています。

・個数制限なしナップサック問題


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

dp = [0] * (W + 1)

for i in range(n):
    for j in range(w[i], W + 1):  # ここだけ違う
        dp[j] = max(dp[j], dp[j - w[i]] + v[i])

print(dp[W])

一次元にして書いてみると、一つ目の01ナップサック問題のコードとの違いはループの向きだけということになります。

例題 2-3-4 01ナップサック問題その2

01ナップサック問題その2


n = int(input())
w = []
v = []
for i in range(n):
    w_, v_ = map(int, input().split())
    w.append(w_)
    v.append(v_)
W = int(input())

dp = [[float("inf")] * (n * 101) for i in range(n + 1)]  # 101 は v_i の制約から

dp[0][0] = 0

for i in range(n):
    for j in range(n * 100):
        if j < v[i]:
            dp[i + 1][j] = dp[i][j]
        else:
            dp[i + 1][j] = min(dp[i][j], dp[i][j - v[i]] + w[i])

res = 0
for i in range(n * 100):
    if dp[n][i] <= W:
        res = i

print(res)

例題 2-3-5 個数制限付き部分和問題

個数制限付き部分和問題

Yes か No かを判定する問題なので、まずは DP 配列の中に bool 値を入れて書いてみましょう。


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

dp = [[False] * (K + 1) for i in range(n + 1)]

# dp[i + 1][j] := i 番目までで j が作れるか
dp[0][0] = True
for i in range(n):
    for j in range(K + 1):
        k = 0
        while k <= m[i] and k * a[i] <= j:
            dp[i + 1][j] |= dp[i][j - k * a[i]]
            k += 1

if dp[n][K]:
    print("Yes")
else:
    print("No")

このアルゴリズムの計算量は $O(K \sum_{i}m_i)$ となり、これでは不十分です。実は DP 配列に bool 値を入れるよりも、同じ計算量でより多くの情報を持たせることができ、そうすることで全体の計算量を落とすことができます。この問題の場合、作れるかどうかだけではなく、作れる場合にあとどれだけ $a_i$ が残っているかを持たせることで計算量を $O(nK)$ まで落とすことができます。


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

dp = [[-1] * (K + 1) for i in range(n + 1)]

# dp[i + 1][j] := i 番目までで j を作る際余る最大の i 番目の個数(作れない場合は -1)
dp[0][0] = 0
for i in range(n):
    for j in range(K + 1):
        if dp[i][j] >= 0:
            dp[i + 1][j] = m[i]
        elif j < a[i] or dp[i + 1][j - a[i]] <= 0:
            dp[i + 1][j] = -1
        else:
            dp[i + 1][j] = dp[i + 1][j - a[i]] - 1

if dp[n][K] >= 0:
    print("Yes")
else:
    print("No")

また、配列の再利用を行って書くと以下のようになります。


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

dp = [-1] * (K + 1)

dp[0] = 0
for i in range(n):
    for j in range(K + 1):
        if dp[j] >= 0:
            dp[j] = m[i]
        elif j < a[i] or dp[j - a[i]] <= 0:
            dp[j] = -1
        else:
            dp[j] = dp[j - a[i]] - 1

if dp[K] >= 0:
    print("Yes")
else:
    print("No")

Maximum-Cup 2018 D Many Go Round

計算量は問題ないですが、いつものように Python だと TLE してしまうので PyPy で提出しましょう。また、4行目の 10 ** 10 ですが、普段は float("inf") を使っていますが MLE してしまうため今回は 10 ** 10 にしました(最近の問題であればメモリ制限が大きいため、ここで float("inf") を使っても MLE にはなりません)。MLE してしまうなどの PyPy を使う際の注意点はこちらのjuppyさんの記事に書いてあります。


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

dp = [[10 ** 10] * m for i in range(n + 1)]

# dp[i + 1][j] := i 番目のタンクまでで番号 j の休憩所にいるための最小周回数
dp[0][0] = 1
for i in range(n):
    for j in range(m):
        dp[i + 1][(j + a[i]) % m] = min(dp[i][(j + a[i]) % m], dp[i][j] + (j + a[i]) // m)

if dp[n][l] <= x:
    print("Yes")
else:
    print("No")

例題 2-3-6 最長増加部分列問題

最長増加部分列問題


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

dp = [0] * n

# dp[i] := 最後が a_i であるような最長の増加部分列の長さ
res = 0
for i in range(n):
    dp[i] = 1
    for j in range(0, i):
        if a[j] < a[i]:
            dp[i] = max(dp[i], dp[j] + 1)
    res = max(res, dp[i])

print(res)

この解き方だと計算量は $O(n^2)$ です。
次は別の漸化式を作ってみます。


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

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

# dp[i] := 長さが i + 1 であるような増加部分列における最終要素の最小値(存在しない場合は INF)
for j in range(n):
    for i in range(n):
        if i == 0 or dp[i - 1] < a[j]:
            dp[i] = min(dp[i], a[j])

res = 0
for i in dp:
    if i < float("inf"):
        res = i

print(res)

この解き方も計算量は $O(n^2)$ ですが、こちらは以下のように二分探索を用いることで $O(n \log n)$ で解くことができます。


from bisect import bisect_left

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

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

# dp[i] := 長さが i + 1 であるような増加部分列における最終要素の最小値(存在しない場合は INF)
for i in range(n):
    dp[bisect_left(dp, a[i])] = a[i]

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

bisect は 二分探索を行うための Python の標準ライブラリです。詳しくは「Python標準ライブラリ:順序維持のbisect」に書かれています。

ABC 006 D トランプ挿入ソート

昇順になっていないカードのみ並び替えれば良いので、すでに昇順になっている枚数を数えます。昇順になっている枚数 = LIS なので、上のコードほぼそのままです。


from bisect import bisect_left

n = int(input())
c = [int(input()) for i in range(n)]

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

# dp[i] := 長さが i + 1 であるような増加部分列における最終要素の最小値(存在しない場合は INF)
for i in range(n):
    dp[bisect_left(dp, c[i])] = c[i]

# 昇順になっていない要素の個数
print(n - bisect_left(dp, float("inf")))

ABC 038 D プレゼント

区間スケジューリング問題と LIS の組み合わせです。二次元になっているのが特徴です。二次元の場合は片方は普通にソートして、もう片方を LIS で求めると上手くいきます。


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

# h_i は既に昇順なので、あとは w_i の LIS を求めれば良い

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")))

TDPC K ターゲット

一つ前のプレゼントと考え方は同じです。


from bisect import bisect_left
from operator import itemgetter

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

# 円の左端と右端を見れば内部にあるかがわかる
xlr = []
for i in c:
    xlr.append((i[0] - i[1], i[0] + i[1]))

# 条件を満たすために左端は降順にソートする
# 境界が重ならないように右端も降順にソートしておく
xlr = sorted(xlr, key=itemgetter(1), reverse=True)  # 右端
xlr = sorted(xlr, key=itemgetter(0), reverse=True)  # 左端

# 右端も条件を満たすように LIS を求める

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

for i in range(n):
    dp[bisect_left(dp, xlr[i][1])] = xlr[i][1]

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

例題 2-3-7 分割数

分割数


n = int(input())
m = int(input())
M = int(input())

dp = [[0] * (n + 1) for i in range(m + 1)]

# dp[i][j] := j の i 分割の総数

dp[0][0] = 1
for i in range(m + 1):
    for j in range(n + 1):
        if j - i >= 0:
            dp[i][j] = (dp[i - 1][j] + dp[i][j - i]) % M
        else:
            dp[i][j] = dp[i - 1][j]

print(dp[m][n])

例題 2-3-8 重複組合せ

重複組合せ


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

dp = [[0] * (m + 1) for i in range(n + 1)]

# dp[i + 1][j] := i 番目までの品物から j 個選ぶ組み合わせの総数

# 1 つも選ばない方法は常に一通り
for i in range(n + 1):
    dp[i][0] = 1

for i in range(n):
    for j in range(1, m + 1):
        if j - 1 - a[i] >= 0:
            # 引き算が含まれる場合には負の数にならないように注意する
            dp[i + 1][j] = (dp[i + 1][j - 1] + dp[i][j] - dp[i][j - 1 - a[i]] + M) % M
        else:
            dp[i + 1][j] = (dp[i + 1][j - 1] + dp[i][j]) % M

print(dp[n][m])
214
202
6

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
214
202