累積最大値（旅館の工事）

問題

ある旅館では、$1$ 号室から $ N $ 号室まで $ N $ 個の部屋があります。

$ i $ 号室は$ A_i $ 人部屋となっています。

この旅館では$ D $日間の工事が行われる予定で、$d$日目は$ L_d $号室から$ R_d $号室までを使うことができません。

$ d=1,2,…,D $のそれぞれの工事日について、稼働できる部屋のうち最も大きい部屋が何人部屋かを求めて下さい。

[制約]
🔹$ 3 \leqq N \leqq 100000 $
🔹$ 1 \leqq D \leqq 100000 $
🔹$ 1 \leqq A_i \leqq 100 $
🔹$ 2 \leqq L_d \leqq R_d \leqq N - 1 $

解き方・ソースコード

まずは、for文を使って泊まれる人数の最大値を計算する方法が思い浮かびますが、計算量が$ O(ND) $ となり実行するのに時間がかかってしまいます。

計算量を減らすために、$L$ 号室から $R$ 号室まで以外の部屋が使えるとき「最も大きい部屋は何人部屋か」を次の式で計算します。

$$ max(1～(L-1)号室の最大人数, (R+1)～N号室の最大人数) $$

１号室から$i$号室までの最大人数$ P_i $、および$i$号室から$N$号室までの最大人数$Q_i$を求めると、$d$日目の答えは$ max(Q_{L_d - 1},Q_{R_d + 1}) $で計算することができます。

次に、$P_i,Q_i$を求めるために、累積和と同じような方法を使うことができます。

$P_i$については先頭要素$A_i$から累積的に最大値を計算し、$Q_i$については最終要素$A_N$から累積的に最大値を計算しておくと、計算量$O(N)$でそれぞれの値を求められます。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
#--------- 入力例 1 ---------
N = 7                       # 部屋の数
A = [1, 2, 5, 5, 2, 3, 1]   # 部屋ごとの宿泊できる人数
D = 2                       # 工事する日数
L = [3, 5]                  # 工事する部屋の番号（開始位置）
R = [4, 6]                  # 工事する部屋の番号（終了位置）
#----------------------------
# P[i] を求める（配列は0からスタート）。最初の部屋から最後の部屋へ向かっての累積[最大値]
P = [ None ] * N
P[0] = A[0]
for i in range(1, N):
    P[i] = max(P[i - 1], A[i])

# Q[i] を求める（配列は0からスタート）。最後の部屋から最初の部屋へ向かっての累積[最大値]
Q = [ None ] * N
Q[N - 1] = A[N - 1]
for i in reversed(range(0, N - 1)):
    Q[i] = max(Q[i + 1], A[i])

# それぞれの日について答えを算出（配列は0からスタート）
for d in range(D):
    print('工事{}日目に使える最も大きい部屋数は{}人部屋'.format(d + 1, max(P[(L[d] - 1) - 1], Q[(R[d] + 1) - 1])))

[実行結果（入力例１）]

工事1日目に使える最も大きい部屋数は3人部屋

工事2日目に使える最も大きい部屋数は5人部屋

工事日ごとに使える部屋のうち、最も大きい部屋が何人部屋かを求めることができました。