Pythonでリストのインデックスkに配置できる最大値を求めるプログラム

この記事では、「サイズn・合計total・隣接要素の差が1以内」という条件を満たすリストの中で、指定したインデックスkに配置できる最大の値をPythonで求める方法を解説します。

問題の概要

3つの値 n(リストのサイズ)、total(要素の合計)、k(インデックス)が与えられます。ここで、次の条件をすべて満たすサイズnのリストを考えます。

• リストの全要素の合計がtotalと一致する
• 隣り合う任意の2要素の絶対差が1以下である

このようなリストの中で、インデックスkの位置に現れうる最大の値を求めるのが目的です。

たとえば、入力が n = 5、total = 15、k = 3 の場合、出力は 4 になります。これは、[3, 2, 3, 4, 3] のようなリストが条件を満たしており、インデックス3(0始まり)には最大で4を配置できるためです。

アプローチ(解き方)

この問題は、候補となる値xを1ずつ増やしながら「インデックスkの値がxであるときに必要な最小の合計」を計算し、その合計が初めてtotalを超える直前のxを答えとする方法で解けます。

隣接差1以下という制約のもとで合計を最小にするには、インデックスkから左右に離れるにつれて要素を1ずつ減らしていく構成が最適です。具体的な手順は以下の通りです。

• x := 0 で初期化する
• 以下を繰り返し実行する:
  • a := k + 1(インデックスkを含む左側の要素数)
  • s := (x + x − a + 1) × a ÷ 2 の小数点以下切り捨て(左側の最小合計)
  • a := n − k(右側の要素数)
  • s := s + (x + x − a + 1) × a ÷ 2 の小数点以下切り捨て(右側の最小合計)
  • s := s − x(インデックスkを左右両方で数えているため、1回分を引く)
  • s > total ならループを抜ける
  • x := x + 1
• x − 1 を返す

この探索は線形に行っていますが、必要な合計sはxに対して単調に増加するため、二分探索を使えばさらに高速化できます。

コード例

理解を深めるために、以下の実装を見てみましょう。

def solve(n, total, k):
    x = 0
    while 1:
        a = k + 1
        s = (x + x - a + 1) * a // 2
        a = n - k
        s += (x + x - a + 1) * a // 2
        s -= x
        if s > total:
            break
        x += 1
    return x - 1

n = 5
total = 15
k = 3
print(solve(n, total, k))

入力

5, 15, 3

出力

4