Pythonで制約付きの建物の最大高さを求めるプログラム

問題の概要

整数 n と制約リスト restrictions が与えられたとします。私たちは都市に n 棟の新しい建物を一列に建てようとしていますが、高さに関するいくつかの制限があります。建物には左から順に 1 から n までの番号が付けられており、各制約は restrictions[i] = (id_i, max_height_i) の形式で表され、「id_i 番の建物の高さは max_height_i 以下でなければならない」ことを意味します。

建物の高さに関する都市の規則は以下のとおりです。

• 各建物の高さは 0 以上でなければなりません。
• 1 番の建物(最初の建物)の高さは必ず 0 です。
• 隣接する 2 つの建物の高さの差は 1 を超えてはいけません。

これらの条件のもとで、最も高い建物が取りうる最大の高さを求める必要があります。

入力例と出力

たとえば、n = 5、restrictions = [[2,1],[4,3]] の場合、出力は 4 になります。このとき最適な高さの並びは [0, 1, 2, 3, 4] となり、2 番の建物(高さ 1 ≤ 1)と 4 番の建物(高さ 3 ≤ 3)の制約をすべて満たしながら、5 番の建物の高さ 4 が最大値となります。

解法のアプローチ

この問題は、制約を左右両方向に伝播させて実現可能な高さに絞り込み、その上で各区間ごとに到達できる最大高さを計算することで解けます。手順は以下のとおりです。

1. 制約リストが空の場合は n - 1 を返します。1 番の建物の高さが 0 なので、以降は毎回 +1 ずつ増やせて、n 番の建物の高さ n-1 が最大になるためです。
2. 制約リストを id の昇順にソートします(resi)。
3. 現在の高さ k = 0、現在位置 idx = 1 で初期化します。
4. 左から右への伝播： 各制約について re[1] を「直前の制約地点から到達できる高さ k + re[0] - idx」と比較し、小さい方に更新します。隣接する建物の高さ差が 1 以内という規則のため、距離的に間に合わない高い制約は実現できないからです。更新後、k と idx を現在の制約の値に置き換えます。
5. リスト末尾の制約の高さと id を k、idx に設定し、リストを反転します。
6. 右から左への伝播： 反転したリストに対して同様の走査を行い、re[1] を「k - re[0] + idx」と比較して小さい方に更新します。完了したら再度反転して元の順序に戻します。
7. f = 0(現在の実現可能な高さ)、idx = 1、res = 0(区間内の最大高さ)で初期化します。
8. 各制約について、その地点で実現可能な高さ ff = min(f + re[0] - idx, re[1]) を求め、区間内の最大高さを res = max(res, (re[0] - idx + f + ff) // 2) として更新します。これは、区間の長さと両端の高さから、途中で折り返す形が最適になることに基づいています。
9. 最後に、「最後の制約地点から n 番の建物まで到達できる高さ f + n - idx」と「各区間で求めた最大値 res」の大きい方を返します。

Pythonでの実装例

それでは、上記のアルゴリズムを実装したコードを見てみましょう。

def solve(n, restrictions):
    if not restrictions:
        return n - 1

    resi = sorted(restrictions, key=lambda x: x[0])

    # 左から右へ制約を伝播
    k = 0
    idx = 1
    for re in resi:
        re[1] = min(re[1], k + re[0] - idx)
        k = re[1]
        idx = re[0]

    # 右から左へ制約を伝播
    k = resi[-1][1]
    idx = resi[-1][0]
    resi.reverse()
    for re in resi[1:]:
        re[1] = min(re[1], k - re[0] + idx)
        k = re[1]
        idx = re[0]
    resi.reverse()

    # 各区間の最大高さを計算
    f = 0
    idx = 1
    res = 0
    for re in resi:
        ff = min(f + re[0] - idx, re[1])
        res = max(res, (re[0] - idx + f + ff) // 2)
        idx = re[0]
        f = ff

    return max(f + n - idx, res)

n = 5
restrictions = [[2, 1], [4, 3]]
print(solve(n, restrictions))

入力

5, [[2,1],[4,3]]

出力

4

まとめ

このアルゴリズムは、制約をソートして左右双方向に伝播させることで各建物の実現可能な高さの上限を確定し、隣接する制約間の区間ごとに到達可能な最大高さを求めます。計算量は制約数を m とすると O(m log m)(ソート部分が支配的)であり、効率的に答えを得られるのが特徴です。