蟻本をpythonで（chapter3 中級編~）

#はじめに
プログラミングコンテストチャレンジブック（通称:蟻本）の問題をpythonで解いていきます。
類題も交えながら少しずつ更新していく予定です。

初級編については@sabaさんの記事を参照するのが良いと思います。
https://qiita.com/saba/items/affc94740aff117d2ca9

#3-1 値の探索だけじゃない！"二分探索"
###p128 lower_bound

N = 5
A = [2, 3, 3, 5, 6]
K = 3

# K以上であるAiのうち最も左にあるもののインデックス
print(bisect_left(A, K))

###p129 cable master

N = 4
K = 11
L = [8.02, 7.43, 4.57, 5.39]

def f(target):
    cnt = 0
    for i in range(N):
        cnt += math.floor(L[i] / target) #各紐から長さtargetの紐がいくつ取れるか
    if cnt >= K: #K本以上取れれば
        return True
    else:
        return False

ng = 0
ok = 10 ** 9 + 1
for i in range(100):
    mid = (ng + ok) / 2
    if f(mid):
        ng = mid
    else:
        ok = mid
print(math.floor(ok * 100) / 100) # 小数点２位まで求める

###p131 aggressive cows

N = 5
M = 3
X = [1, 2, 8, 4, 9]
X.sort()

def c(d):
    last = 0 # １つ目の小屋を置く
    # M - 1回現在の場所からd以上離れた場所に小屋を置けるか
    for _ in range(M - 1):
        crt = last + 1 # crtを進めて次に小屋を置く場所を探す
        while crt < N and X[crt] - X[last] < d:
            crt += 1 

        if crt == N: # 次に小屋を置く場所が見つからなければ
            return False

        last = crt # 次の小屋を置く

    return True # M - 1回小屋を置ければ

ok = -1
ng = 10 ** 9 + 1

while abs(ok - ng) > 1:
    mid = (ok + ng) // 2
    if c(mid):
        ok = mid
    else:
        ng = mid
print(ok)

###p132 平均最大化

N = 3
K = 2
W = [2, 5, 2]
V = [2, 3, 1]

def c2(x, w, v):
    # 「単位重さあたりの価値がx以上になる」を満たすよう選んだ商品の集合を
    # S([w1, v1], [w2, v2]...[wk, vk])とする。この時
    # sum(v1, v2...vk) / sum(w1, w2...wk) >= x より
    # sum(v1, v2...vk) - x * sum(w1, w2...wk) >= 0
    # sum(v1 - x * w1, v2 - x * w2, ...vk - x * wk) >= 0

    # v[i] - x * w[i]を大きい順にk個取ったときそれが0以上になればいい
    cnt = 0
    items = [v[i] - x * w[i] for i in range(N)]
    items.sort(reverse = True)
    for i in range(K):
        cnt += items[i]
    return cnt >= 0

ok = 0
ng = 10 ** 9 + 1

for i in range(100):
    mid = (ng + ok) / 2
    if c2(mid, W, V):
        ok = mid
    else:
        ng = mid
print(ok)

#3-2 厳選！頻出テクニック(1)

p135 subsequence

N = 10
S = 15
A = [5, 1, 3, 5, 10, 7, 4, 9, 2, 8]
right = 0
total = 0
ans = float('inf')

for left in range(N):
    while right < N and total < S:
        # 条件を満たすまでrightを進める total >= Sになった時点で打ち切る

        # left, right, total = 0
        # total < SなのでA[0]をtotalに足してrightを１つ増やす
        # left, right, total = 0, 1, 5
        # ...
        # left, right, total = 0, 4, 14
        # total < SなのでA[4]をtotalに足して(この時点でtotal < Sになる）rightを１つ増やす
        # left, right, total = 0, 5, 24
        # total >= SなのでA[5]を足さない

        total += A[right]
        right += 1

    # もしtotal < Sのままrightが右端まで達したら
    if total < S:
        continue
 
    # left, right = 0, 5なら閉区間[0, 4]の総和がS以上になる
    # → 半開区間[0, 5)の総和がS以上になる
    ans = min(ans, right - left)

    if left == right:
        right += 1
    total -= A[left]

    # print(left, right)

print(ans)

p137 Jessica's Reading Problem

P = 5
A = [1, 8, 8, 8, 1]
dict = {}
for i in A:
    dict[i] = 0

l = 0
cnt = 0
ans = float('inf')
# 右を一つ進めて左をできる限り削る
for r in range(P):
    # 新しく一つ加える
    if dict[A[r]] == 0: # 新しい要素が来たなら
        cnt += 1
    dict[A[r]] += 1

    # 全ての要素を含みながら左を削っていく
    while True:
        # もし要素が１つしか無かったら削れない
        if dict[A[l]] == 1:
            break
        # 要素が２つ以上あれば削る
        dict[A[l]] -= 1
        l += 1

    if cnt < len(dict.items()):
        continue

    ans = min(ans, r - l + 1)
print(ans)

p138 Face The Right Way

N = 7
S = 'BBFBFBB'

def judge(k):
    imos = [0] * N
    # 最初の一回をひっくり返すかどうか
    if S[0] == 'B':
        imos[0] = 1
    # 前から順に牛が後ろを向いているならひっくり返す
    # k = 4なら
    # BBFBFBB
    # [  ]
    #  [  ]
    #   [  ]
    #    [  ]
    # 最大4回ひっくり返す
    for i in range(1, N - k + 1):
        if i < k:
            rev = imos[i - 1]
        else:
            rev = imos[i - 1] - imos[i - k]
        if (S[i] == 'B') ^ (rev % 2):
            imos[i] += 1
        imos[i] += imos[i - 1]

    # 残りの牛が前を向いているか調べる
    for i in range(N - k + 1, N):
        if i < k:
            rev = imos[N - k]
        else:
            rev = imos[N - k] - imos[i - k]
        if (S[i] == 'B') ^ (rev % 2):
            return float('inf')

    return imos[N - k]

K = 0
M = float('inf')
for i in range(1, N + 1):
    opt = judge(i)
    if opt < M:
        M = opt
        K = i
print(K, M)

p141 Fliptile

M = 4
N = 4
maze = [
[1, 0, 0, 1],
[0, 1, 1, 0],
[0, 1, 1, 0],
[1, 0, 0, 1]
]

def changer(o, f):
    for i in range(1, M):
        for j in range(N):
            # 一つ上のmazeとflipを調べる
            # mazeが黒 and flipが偶数回ひっくりかえる or mazeが白 and flipが奇数回ひっくりかえるなら
            if (maze[i - 1][j] == 1) ^ (f[i - 1][j] % 2):
                o[i][j] += 1
                f[i][j] += 1
                f[i - 1][j] += 1
                if j >= 1:
                    f[i][j - 1] += 1
                if j < N - 1:
                    f[i][j + 1] += 1
                if i < M - 1:
                    f[i + 1][j] += 1

    return o, f

# １行目について全探索
# １行目が定まると２行目以降が一意に定まる
# １行目について全探索:O(2 ** N) ２行目以降を埋めていく:O(MN)かかるので
# 合計の計算量はO(MN(2 ** N))になる
for bit in range(1 << N):
    opt = [[0] * N for i in range(M)]
    flip = [[0] * N for i in range(M)]
    for i in range(N):
        if bit & (1 << i):
            opt[0][i] += 1
            flip[0][i] += 1
            if i >= 1:
                flip[0][i - 1] += 1
            if i < N - 1:
                flip[0][i + 1] += 1

    opt, flip = changer(opt, flip)

    # 最後の列について判定
    for j in range(N):
        # 違うなら
        if (maze[M - 1][j] == 1) ^ (flip[M - 1][j] % 2):
            break
    else:
        for i in opt:
            print(*i)

p145 physics experiment

# antsと同じく玉同士はすり抜けるものとして考える
N, H, R, T = 2, 10, 10, 100
g = 10

def judge(time):
    if time < 0:
        return H
    t = math.sqrt(2 * H / 10)
    k = int(time / t)
    if k % 2 == 0:
        d = time - k * t
        return H - g * d * d / 2
    else:
        d = k * t + t - time
        return H - g * d * d / 2

ans = []
for i in range(N):
    # 一秒ごとにボールを落下させる
    ans.append(judge(T - i))
# ボールは互いにすり抜けるものと考えて良い
ans.sort()
for i in range(N):
    # RはセンチメートルだがHはメートルなのに注意
    print(ans[i] + (2 * R * i / 100))

p147 4 values whose sum is 0

# 二分探索
def binary_search_loop(data, target):
    imin = 0
    imax = len(data) - 1
    while imin <= imax:
        imid = imin + (imax - imin) // 2
        if target == data[imid]:
            return imid
        elif target < data[imid]:
            imax = imid - 1
        else:
            imin = imid + 1
    return False

N = 6
A = [-45, -41, -36, -36, 26, -32]
B = [22, -27, 53, 30, -38, -54]
C = [42, 56, -37, 75, -10, -6]
D = [-16, 30, 77, -46, 62, 45]

re_A = []
re_C = []
# Ai + Bj
for i in range(N):
    for j in range(N):
        re_A.append(A[i] + B[j])
re_A.sort()
# Ci + Dj
for i in range(N):
    for j in range(N):
        re_C.append(C[i] + D[j])
re_C.sort()

ans = 0
for i in re_A:
    # ind1:合計で0以上になるもののうち一番左のインデックス
    # ind2:合計で1以上になるもののうち一番左のインデックス
    #             ind1       ind2
    # [...-2, -1, (0), 0, 0, (2), 2, 3]
    ind1 = bisect_left(re_C, 0 - i)
    ind2 = bisect_left(re_C, 0 - i + 1)
    ans += ind2 - ind1
print(ans)

p148 巨大ナップサック

N = 4
w = [2, 1, 3, 2]
v = [3, 2, 4, 2]
W = 5

# 半分全列挙 + 二分探索
def re_list(weight, value):
    fore_list = []
    # まず全通り組み合わせる
    for bit in range(1 << len(weight)):
        wei = 0
        val = 0
        for i in range(len(weight)):
            if bit & (1 << i):
                wei += weight[i]
                val += value[i]
        fore_list.append([wei, val])
    fore_list.sort()

    # リスト再作成
    # 明らかにいらないものは消している
    # ABC032 D - ナップサック問題の解説参照
    alta_w = []
    alta_v = []
    now = -1
    for i in fore_list:
        if now < i[1]:
            now = i[1]
            alta_w.append(i[0])
            alta_v.append(i[1])
    return alta_w, alta_v

def half_knapsack(N, limit, weight, value):
    # 半分全列挙
    fore_w, fore_v = re_list(weight[:N // 2], value[:N // 2])
    back_w, back_v = re_list(weight[N // 2:], value[N // 2:])

    ans = 0
    for b_w, b_v in zip(back_w, back_v):
        if b_w > limit:
            continue

        opt = b_v
        index = bisect_right(fore_w, limit - b_w)
        if index > 0:
            opt += fore_v[index - 1]
        ans = max(ans, opt)

    return ans
print(half_knapsack(N, W, w, v))

###p150 領域の個数

# 与えられたx, yの上限が小さいため今回の入力例では全く圧縮されないんです
W, H, N = 10, 10, 5
# (x1, y1)から(x2, y2)まで線を引く
X1 = [i - 1 for i in [1, 1, 4, 9, 10]]
X2 = [i - 1 for i in [6, 10, 4, 9, 10]]
Y1 = [i - 1 for i in [4, 8, 1, 1, 6]]
Y2 = [i - 1 for i in [4, 8, 10, 5, 10]]

# 縦
row = set()
r_index = {}
# 横
col = set()
c_index = {}
# 変換
for x, y in zip(X1 + X2, Y1 + Y2):
    # どの横列が必要か
    row.add(y)
    if y > 0:
        row.add(y - 1)
    if y < H - 1:
        row.add(y + 1)

    # どの縦列が必要か
    col.add(x)
    if x > 0:
        col.add(x - 1)
    if x < W - 1:
        col.add(x + 1)

# 圧縮後のどの座標になるか
H = len(row)
row = sorted(list(row))
for i in range(H):
    r_index[row[i]] = i

W = len(col)
col = sorted(list(col))
for i in range(W):
    c_index[col[i]] = i

X1 = [c_index[i] for i in X1]
X2 = [c_index[i] for i in X2]
Y1 = [r_index[i] for i in Y1]
Y2 = [r_index[i] for i in Y2]

# print(X1, Y1, X2, Y2)
# X1: [0, 0, 3, 8, 9]
# Y1: [3, 7, 0, 0, 5]

# X2: [5, 9, 3, 8, 9]
# Y2: [3, 7, 9, 4, 9]

# あとは普通に解く
maze = [[0] * W for i in range(H)]
# 線を引く
for x1, y1, x2, y2 in zip(X1, Y1, X2, Y2):
    # 縦線
    if x1 == x2:
        if y1 > y2:
            y1, y2 = y2, y1
        for i in range(y1, y2 + 1):
            maze[i][x1] = 1
    # 横線
    else:
        if x1 > x2:
            x1, x2 = x2, x1
        for i in range(x1, x2 + 1):
            maze[y1][i] = 1

# p35 lake counting する
dx = [0, 1, 0, -1]
dy = [1, 0, -1, 0]

def dfs(y, x):
    maze[y][x] = 1
    for i in range(4):
        ny = y + dy[i]
        nx = x + dx[i]
        if 0 <= ny < H and 0 <= nx < W and maze[ny][nx] == 0:
            dfs(ny, nx)

ans = 0
for y in range(H):
    for x in range(W):
        if maze[y][x] == 0:
            dfs(y, x)
            ans += 1
print(ans)

3-3 さまざまなデータ構造を知ろう