Pythonで配列の要素を削除して得られる最大ポイントを求める方法

問題の概要

N個の要素を持つ配列 A と、2つの整数 l・r が与えられます(要素の値は 1 ≤ ax ≤ 10^5、かつ 1 ≤ l ≤ r ≤ N を満たします)。配列から任意の要素 ax を1つ取り除くと、それと同時に「ax+1、ax+2、…、ax+R」および「ax−1、ax−2、…、ax−L」に等しい値をもつすべての要素も配列から削除されます。この操作を行うと ax ポイントを獲得できます。配列のすべての要素を削除し終えたとき、獲得できる合計ポイントの最大値を求めるのが目的です。

たとえば、入力が A = [2,4,3,10,5]、l = 1、r = 2 の場合、出力は 18 になります。

解法のアプローチ

この問題は、各値の出現回数をあらかじめ数えておき、値の昇順に動的計画法(DP)で「その値を取るか・取らないか」を決めていくことで効率的に解けます。res[num] を「値 num 以下の要素だけを処理したときに得られる最大ポイント」と定義すると、次の遷移式が成り立ちます。

res[num] = max(res[num−1], num × count_list[num] + res[k])(ただし k = max(num − left − 1, 0))

具体的な手順は以下の通りです。

• n := 配列のサイズ

• max_val := 0

• i を 0 〜 n−1 でループし、max_val := max(max_val, array[i]) として配列内の最大値を求める

• count_list := サイズ (max_val + 1) の配列を 0 で初期化

• 再度 i をループし、count_list[array[i]] を加算して各値の出現回数を記録する

• res := サイズ (max_val + 1) のDP配列を 0 で初期化(res[0] := 0)

• left := min(left, right)(削除範囲の判定には小さい方の幅を使用)

• num を 1 〜 max_val でループ:

  • k := max(num − left − 1, 0)

  • res[num] := max(res[num − 1], num × count_list[num] + res[k])

• res[max_val] を返す

ここで res[num − 1] は「値 num を取らない場合」の最適解、num × count_list[num] + res[k] は「値 num の要素をすべて取ってポイントを獲得し、削除範囲と重なる小さい側の値(k 以下)まで遡った最適解に足し合わせる場合」を表しています。この遷移により、削除操作による干渉を考慮しながら全体の最大ポイントを求められます。

実装例

以下がPythonでの実装例です。

def get_max_cost(array, left, right) :
    n = len(array)
    max_val = 0
    for i in range(n) :
        max_val = max(max_val, array[i])
    count_list = [0] * (max_val + 1)
    for i in range(n) :
        count_list[array[i]] += 1
    res = [0] * (max_val + 1)
    res[0] = 0
    left = min(left, right)
    for num in range(1, max_val + 1) :
        k = max(num - left - 1, 0)
        res[num] = max(res[num - 1], num * count_list[num] + res[k])
    return res[max_val]

array = [2,4,3,10,5]
left = 1
right = 2
print(get_max_cost(array, left, right))

入力

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

出力

18

計算量

配列の走査に O(n)、DPの計算に O(max_val) かかるため、全体の時間計算量は O(n + max_val)、空間計算量は O(max_val) となります。要素の値が最大 10^5 であることを考えると、十分に高速で実用的な解法といえます。