トップ > AtCoder > ABC236E - Average and Median

ABC236E - Average and Median

AtCoder AtCoder-Beginner Contest Python DP 二分探索 二分探索-平均値 二分探索-中央値

考え方

公式解説がわかりやすい．二分探索＋DP．
Editorial - AtCoder Beginner Contest 236

平均値

\begin{aligned}
& \frac{1}{n} \sum_{i=1}^{n} x_{i} \geq K \\
& \Leftrightarrow \sum_{i=1}^{n} (x_{i} -K) \geq 0
\end{aligned}
と読み替える．


中央値

\begin{aligned}
y_{i} =
\begin{cases}
\, 1 & (x_{i} \geq K)\\
\, -1 & (x_{i} < K)
\end{cases}
\end{aligned}
とする．

「(中央値) $\geq K$」は，$K$以上の値が半分より多い (※)（$\lceil n/2 \rceil $個より多い）ことと同値であるから，

\begin{aligned}
& \mathrm{median\,} (\{x_{1}, x_{2},...,x_{n} \}) \geq K \\
& \Leftrightarrow \sum_{i=1}^{n} y_{i} > 0
\end{aligned}
と読み替えられる．

(※) 例えば，$\{1,1,2,2\}$の中央値は$1$．

二分探索＋DP

これまでの内容を使って，「平均値 / 中央値 が$K$以上になるか？」を二分探索する．その際，「平均値 / 中央値 の最大値」を求めるのには，次のDPを使う．ただし，平均値の計算では$B_{k} = x_{k} -K$とし，中央値の計算では$B_{k} = y_{k}$とする．

  • $\mathrm{used\,}[i] = i$番目まで決定する & $i$番目を使う場合の $\sum_{k} B_{k}$の最大値
  • $\mathrm{unused\,}[i] = i$番目まで決定する & $i$番目を使わない場合の $\sum_{k} B_{k}$の最大値

遷移は

\begin{aligned}
\mathrm{used\,}[i+1]
&=\max(\mathrm{used\,}[i], \mathrm{unused\,}[i]) + B_{i+1} \\
\mathrm{unused\,}[i+1]
&=\mathrm{used\,}[i]
\end{aligned}
となる．

回答例

$2^{60} \sim 10^{18}$だから，誤差okのはず．

N = int(input())
A = list(map(int, input().split()))

# --- average ---
ok, ng = 0, 10 ** 9
for _ in range(60):
  mid = (ok + ng) / 2
  B = [a - mid for a in A]
  used, unused = 0, 0
  for b in B:
    used, unused = max(used, unused) + b, used
  if max(used, unused) >= 0:
    ok = mid
  else:
    ng = mid
print(ok)

# --- median ---
ok, ng = 0, 10 ** 9 + 1
while ng - ok > 1:
  mid = (ok + ng) // 2
  B = [1 if a >= mid else -1 for a in A]
  used, unused = 0, 0
  for b in B:
    used, unused = max(used, unused) + b, used
  if max(used, unused) > 0:
    ok = mid
  else:
    ng = mid
print(ok)

【参考】
Submission #28748466 - AtCoder Beginner Contest 236
【解説 実況】ABC236 E問題【かつっぱ】 - YouTube