平均値を最大化する（条件付き、要計算量削減）【競技プログラミング】

# はじめに
- 平均値の最大化、って単純な感じがするけど、面倒な条件があったり、計算量削減が必要だったり、競技プログラミングの醍醐味が味わえるネタだと思うので、取り上げてみる。

# 参考となる問題
- AtCoderの[ABC236e問題](https://atcoder.jp/contests/abc236/tasks/abc236_e)
- これは、配列からの選択条件が、「2個飛ばし不可」というもの。つまり、要素が5個{0,1,2,3,4}あったとして、{0,2,4}とか、{1,2,4}とかはokだけど、{0,3,4}は「1,2」を飛ばしたから駄目。
- なお、入力は以下の形となる。
```
6 // 配列のサイズ6
2 1 2 1 1 10 // 配列の要素
```

# 条件を見たす解を求める
- まぁ、解法としては、dpしか思いつかないよね。
- 配列dp[i][s]について、インデックスiまで遷移した時点で、s個の要素を選択したときの、平均値の最大値とする。sという情報は、平均値を算出するときに必須。
- また、「2個飛ばし不可」を遷移の式に反映させるため、インデックスiの要素を選んだか(st：selectの意味)、選ばなかったか(ns:not selectの意味)の二区分を用意するため、dp[i][s]の中身をタプル(st:Int,ns:Int)とする。
- dpの遷移式は、次のとおり
    - dp[i+1][s].st については、dp[i][s].stとdp[i][s].nsのうち、大きい方にsを乗じた後にvs[i]を加算し、その後、(s+1)で除した値となる。
    - dp[i+1][s].ns については、「2個飛ばし不可」のルールにより、dp[i][s].st となる。また、「.ns」なのでvs[i]は加算されない。
- 上記に基づき、コードを書くとこんな感じ。
```swift
extension Array {
    init(_ cnt: Int, _ v: Element) {
        self = [Element](repeating: v, count: cnt)
    }
}
let n = Int(readLine()!)!

// 計算しやすくするため、整数値をDouble型で格納
let vs = readLine()!.split(separator:" ").map{Double($0)!}

// dp配列の準備。平均値を格納するため、Double型としている。
var dp = [[(st:Double,ns:Double)]](n+1,[(Double,Double)](n+1,(-1,-1))) // st:セレクト、ns:noセレクト

// 検証用：選んだカード番号を格納 -- selected cards でsc
var sc = [[(st:[Int],ns:[Int])]](n+1,[([Int],[Int])](n+1,([],[])))

// 初期化
dp[0][0].st = 0 // -1を遷移不可フラグとして使っているので、その解除の意味も含む

// 遷移
for i in 0..<n {
    for s in 0...i {
        let ds = Double(s) //計算用にInt->Doubleにしておく
        
        // i+1でnsを選べるのは、直前でstの時だけ
        if dp[i][s].st != -1 { // 遷移前の状態が存在しなければスキップ
            dp[i+1][s].ns = dp[i][s].st
            sc[i+1][s].ns = sc[i][s].st               
        }
        
        // i+1でstを選ぶ場合、直前でns,stとも可能
        if max(dp[i][s].ns,dp[i][s].st) != -1 { // 遷移前の状態が存在しなければスキップ
            if dp[i][s].st > dp[i][s].ns {
                dp[i+1][s+1].st = (ds * dp[i][s].st + vs[i]) / (ds+1)
                sc[i+1][s+1].st = sc[i][s].st + [i]
            } else {
                dp[i+1][s+1].st = (ds * dp[i][s].ns + vs[i]) / (ds+1)
                sc[i+1][s+1].st = sc[i][s].ns + [i]                
            }
        }
    }
}

// 
for (s,dp_n) in dp[n].enumerated() {
    print(s,dp_n) // -- 下記出力①を参照
}

// 検証用の選んだカード番号を表示
for (s,sc_n) in sc[n].enumerated() {
    print(s,sc_n) // -- 下記出力②を参照
}

///出力
①dp[n]の結果(条件を見たす最大平均値)
0 (st: -1.0, ns: -1.0)
1 (st: -1.0, ns: -1.0)
2 (st: -1.0, ns: -1.0)
3 (st: 4.0, ns: 1.6666666666666667)
4 (st: 3.75, ns: 1.5)
5 (st: 3.2, ns: 1.4)
6 (st: 2.8333333333333335, ns: -1.0)

②sc[n]の結果(選んだカード番号の組み合わせ)
0 (st: [], ns: [])
1 (st: [], ns: [])
2 (st: [], ns: [])
3 (st: [1, 3, 5], ns: [0, 2, 4])
4 (st: [0, 2, 3, 5], ns: [0, 1, 2, 4])
5 (st: [0, 1, 2, 3, 5], ns: [0, 1, 2, 3, 4])
6 (st: [0, 1, 2, 3, 4, 5], ns: [])
```
- 配列dpについて説明通りのコードを書いたところ、遷移終了後のdp[n]は出力①のとおり。
- 出力①を見ると、s=0,1,2となるときの値は「-1」であり、「2個飛ばし不可」を満たすことが出来ないことが分かる。また、s=6のnsも「-1」だが、これは最後の一枚を引かない(ns)ときは、6枚引けないため当然である。
- 出力①で「-1」以外の結果を見ると、S=4のstが4で最大値であり、これがコンテストの解となる。
- また、蛇足ではあるが、どのカードを選んだかを確認するため、配列scを用意した。
- 出力②を見ると、S=4のstでは、カード[1,3,5]が選ばれている。
- なお、上記コードは二重forループとなっており、計算量は$O(n^2)$であり、制約は$n≦10^5$となっているから、計算量は最大で$10^{10}$となってしまうので、TLEとなってしまう。

# 計算量の削減
- 計算量の削減と言えば、ニブタンだよね。
- ということで、上記のdpを使ったコードを[solve関数にして、ニブタンで解けないか](https://qiita.com/kei_tnk/items/85d14f6c835b83f1ed07)を考える。
- え？平均値の最大化でニブタン？[最近、考えた気がするなぁ](https://qiita.com/kei_tnk/items/a6b21e709b875b0d1c67)。
- 同じ思想を踏襲して、solve関数の引数をfとし、カードの配列vsから新しい配列 vs_f = vs.map{ $0 - x }を作ることとする。このとき、配列vs_fから、合計が0以上となるようなカードの組み合わせを選択可能なら、solve関数をtrueとして返すこととする。このように、「平均値の問題」を「合計値の問題」に置き換えるのは、定番中の定番みたいだね。
- それでは、このときdpはどのように書けるだろうか？当初のdpでは、平均値を求める必要があったため、それまでに選択した枚数(s)の情報が必要だった。今回は、「平均値」ではなく「合計値」の問題であるため、枚数(s)の情報は不要となり、dpは1次元配列として、以下のように書ける
    - dp[i+1].st = max(dp[i].st,dp[i].ns) + vs_f[i]
    - dp[i+1].ns = dp[i].st
- うーむ、スッキリ！
- さて、そんなdpを使ってsolve関数を作ってみると、以下の通り。
```swift
func solve(_ f:Double) -> Bool{
    let vs_f = vs.map{$0 - f}
    
    // dpの準備
    var dp = [(st:Double,ns:Double)](n+1,(-1,-1)) // st:セレクト、ns:noセレクト
    
    // dpの初期化
    dp[0].st = 0
    
    // dpの遷移
    for i in 0..<n {
        // i+1でnsを選べるのは、直前でstの時だけ
        dp[i+1].ns = dp[i].st
        // i+1でstを選ぶ場合、直前でns,stとも可能
        
        let bigger = dp[i].ns > dp[i].st ? dp[i].ns : dp[i].st // paizaでなぜかmax関数が死ぬため
        dp[i+1].st = bigger + vs_f[i]
        // dp[i+1].st = max(dp[i].ns,dp[i].st) + vs_f[i]
        
    }
    
    return (dp[n].st > dp[n].ns ? dp[n].st >= 0 : dp[n].ns >= 0 ) // paizaでなぜかmax関数が死ぬため
    // return max(dp[n].st,dp[n].ns) >= 0
}
```
- なぜか、max関数が使えなかったから、少し不格好だけど、3項演算子を使用した。
- さて、このsolve関数の計算量はO(n)に下がっているよ。良かったね。

# おまけ〜中央値を最大化する〜
- [ABC236e問題](https://atcoder.jp/contests/abc236/tasks/abc236_e)は、平均値の最大化だけでなく、中央値の最大化も回答しなければならない。よって、こちらも考える。
- 中央値とは、小さい方から数えた順番が真ん中の数値。
    - 5個の数値があったとき、3番目の数値。
    - 8個の数値があったとき、4番目の数値。(5番目ではない)
    - つまり、奇数個数nの場合は、(n+1)/2番目、偶数個数nの場合は、n/2番目の数。
- さて、これも当然、solve関数を使ったニブタンを考える。
- solve関数の引数をx:Intとして、配列fsから選んだ要素の中央値がx以上になるような組み合わせが存在するとにsolve関数がtrueを返せばよい。
- このとき、n枚のカードからs枚のカードを選びx以上のカードの枚数がs/2(sが偶数)もしくは(s+1)/2(sが奇数)以上の枚数であればtrueとなれば良いが、このように条件にsを含む形でdpを用いてsolve関数を作ってしまうと、また、$O(n^2)$の計算量となってしまう。
- よって、カードの配列vsを用いて、vs_x=vs.map{$0>=x ? 1 : -1}という配列を新たに作れば、vs_xから選んだ要素の合計値が0以上となるかどうかの判定に変換できる。
- 上記の考えによりsolve関数を作ると以下の通り
```swift
func solve(_ x:Int) -> Bool{
    let vs_x = vs.map{$0 >= x ? 1 : -1}
    
    // 平均値を格納するため、Double型
    var dp = [(st:Int,ns:Int)](n+1,(-1,-1)) // st:セレクト、ns:noセレクト
    
    // 初期化
    dp[0].st = 0
    // dp[0].ns = 0
    
    // 遷移
    for i in 0..<n {
        // i+1でnsを選べるのは、直前でstの時だけ
        dp[i+1].ns = dp[i].st
        // i+1でstを選ぶ場合、直前でns,stとも可能
        
        let bigger = dp[i].ns > dp[i].st ? dp[i].ns : dp[i].st // paizaでなぜかmax関数が死ぬため
        dp[i+1].st = bigger + vs_x[i]
        // dp[i+1].st = max(dp[i].ns,dp[i].st) + vs_f[i]
        
    }
    
    return (dp[n].st > dp[n].ns ? dp[n].st > 0 : dp[n].ns > 0 ) // paizaでなぜかmax関数が死ぬため
    // return max(dp[n].st,dp[n].ns) > 0
}
```
- 上記も計算量はO(n)となる。

# [ABC236e問題](https://atcoder.jp/contests/abc236/tasks/abc236_e)の解答
- 前項までの平均値と中央値の最大化をガッチャンコしたコードは以下の通り。
```swift
extension Array {
    init(_ cnt: Int, _ v: Element) {
        self = [Element](repeating: v, count: cnt)
    }
}
let n = Int(readLine()!)!
let vs = readLine()!.split(separator:" ").map{Int($0)!}

func solve_ave(_ f:Double) -> Bool{
    let vs_f = vs.map{Double($0) - f}
    
    // dpの準備
    var dp = [(st:Double,ns:Double)](n+1,(-1,-1)) // st:セレクト、ns:noセレクト
    
    // dpの初期化
    dp[0].st = 0
    
    // dpの遷移
    for i in 0..<n {
        // i+1でnsを選べるのは、直前でstの時だけ
        dp[i+1].ns = dp[i].st
        // i+1でstを選ぶ場合、直前でns,stとも可能
        
        let bigger = dp[i].ns > dp[i].st ? dp[i].ns : dp[i].st // paizaでなぜかmax関数が死ぬため
        dp[i+1].st = bigger + vs_f[i]
        // dp[i+1].st = max(dp[i].ns,dp[i].st) + vs_f[i]
        
    }
    
    return (dp[n].st > dp[n].ns ? dp[n].st >= 0 : dp[n].ns >= 0 ) // paizaでなぜかmax関数が死ぬため
    // return max(dp[n].st,dp[n].ns) >= 0
}

func solve_med(_ x:Int) -> Bool{ // 中央値はmedianだからmed
    let vs_x = vs.map{$0 >= x ? 1 : -1}
    
    var dp = [(st:Int,ns:Int)](n+1,(-1,-1)) // st:セレクト、ns:noセレクト
    
    // 初期化
    dp[0].st = 0
    // dp[0].ns = 0
    
    // 遷移
    for i in 0..<n {
        // i+1でnsを選べるのは、直前でstの時だけ
        dp[i+1].ns = dp[i].st
        // i+1でstを選ぶ場合、直前でns,stとも可能
        
        let bigger = dp[i].ns > dp[i].st ? dp[i].ns : dp[i].st // paizaでなぜかmax関数が死ぬため
        dp[i+1].st = bigger + vs_x[i]
        // dp[i+1].st = max(dp[i].ns,dp[i].st) + vs_f[i]
        
    }
    
    return (dp[n].st > dp[n].ns ? dp[n].st > 0 : dp[n].ns > 0 ) // paizaでなぜかmax関数が死ぬため
    // return max(dp[n].st,dp[n].ns) >= 0
}

let f_min:Double = 1 // 検索範囲の最小値
let f_max:Double = 1000000000 // 検索範囲の最大値

var f_ok = f_min // (左側)の開始位置 -- ループ毎に「同位置」or「中間値」へ移動
var f_ng = f_max// （右側）の開始位置 -- ループ毎に「同位置」or「中間値」へ移動

while f_ng - f_ok > 0.001 {
    let mid = (f_ok + f_ng) / 2
    // print(mid,ok,ng,solve(mid)) // 出力を参照
    if solve_ave(mid){
        f_ok = mid
    } else {
        f_ng = mid
    }
}

print(f_ok) // 「平均値の最大化」の答え

let x_min:Int = 1 // 検索範囲の最小値
let x_max:Int = 1000000000 // 検索範囲の最大値

var x_ok = x_min - 1 // (左側)の開始位置 -- ループ毎に「同位置」or「中間値」へ移動
var x_ng = x_max + 1// （右側）の開始位置 -- ループ毎に「同位置」or「中間値」へ移動

while x_ng - x_ok > 1 {
    let mid = (x_ok + x_ng) / 2
    // print(mid,ok,ng,solve(mid)) // 出力を参照
    if solve_med(mid){
        x_ok = mid
    } else {
        x_ng = mid
    }
}

print(x_ok) // 「中央値の最大化」の答え
```

# 最後に
- dpやsolve関数の作り方について、自分で書いてみて、徐々に掴みかけている気がするよ。
- 天才肌の人は、一発で理解するんだろうけど、自分の場合は、反復練習が必要みたいだ。