ABC366 Fについて、自分の解法が公式解説とは違ったので記録しておきます。
問題概要
$N$個の1次関数$f_i(x)=A_ix+B_i$が与えられるので、そのうち相異なる$K$個を選んで合成関数を作り、$x=1$における合成関数の最大値を求めましょう。
-
$1 \le N \le 200,000$
-
$1 \le K \le \min(N,10)$
-
$1 \le A_i,B_i \le 50$
考えたこと
$A_i,B_i$が正であることから、$x$が大きいほど$f_i(x)$の値が大きくなることが解ります。
また$K$の制約が小さいので、選んだ$f_i$を覚えながらDPすれば良さそうなです。
相異なる$N$以下の$k$個の正の整数の集合を$P_k=\set{p_1,…,p_k}$と書くことにします。
この時、$i \in P_k$を満たす$f_i$の合成関数の値の最大値を、$dp[P_k]$で表すことにすると、下記のような漸化式が考えられます。
\displaylines{
dp[\phi]=1\\
dp[P_{k+1}] = \max(f_i(dp[P_k])) (ただしi \notin P_k)}
このままだと、$dp[P_k]$ のパターン数が $ {}_N \mathrm{C} _K $ となるため、TLEです。
高速化の工夫
状態数削減
最適解の$P_k$の要素数は高々$\max(K)=10$です。
また、$x$を固定した時、各$f_i(x)$の値は一位に決まります。
ということは、$dp[P_k]$から$dp[P_{k+1}]$を求めるとき、$f_i$のうち上位から$K$個の中に最適解を構成する$f_i$が含まれているはずです。
事前計算
$x$を固定した時の各$f_i(x)$の値を前計算しておきます。
$B_i \le 50$ より、$x \gt 50$ の時の$f_i$の大小関係は$x=50$の時と同じですので、前計算は$1 \le x \le 50$の範囲で実施すれば十分です。1
$x$の値を固定した時、$f_{p_1}(x) \ge f_{p_2}\ge … \ge f_{p_N}$となるように$p_i$をソートし、ソート結果を$q[x]$と書くことにします。$q[x]$の$i$番目の値が$q[x][i]$です。
また、上述の通り、記憶するのはソート結果のうち上位$K$個で十分です。
ACコード
上記の工夫を実装すると下記のようになります。
Rustだと、BTreemapのキーにBTreeSetを放り込んだりできて楽ですね。
ちなみに実行時間は144ms。少し工夫して、前計算部分を必要になった時だけ実行するようにすると、37msになりました。
#![allow(non_snake_case)]
use proconio::input;
use std::cmp::*;
use std::collections::*;
fn main() {
input! {
N:usize, K:usize, AB:[(i64,i64);N]
}
let mut abi = vec![];
for i in 0..N {
abi.push((AB[i].0, AB[i].1, i));
}
// xを固定した時のf_i(x)のうち、上位K個を事前計算
let mut fp_max = vec![vec![]; 51];
for x in 1..=50 {
abi.sort_by_key(|&(a, b, _)| a * x + b);
for i in 0..K {
fp_max[x as usize].push(abi[N - 1 - i])
}
}
let mut dp = BTreeMap::new();
dp.insert(BTreeSet::new(), 1);
let mut ans = 0;
for _ in 0..K {
let mut np = BTreeMap::new();
for (set_prev, &val_prev) in dp.iter() {
for &(a, b, i) in fp_max[min(val_prev, 50) as usize].iter() {
if set_prev.contains(&i) {
continue;
}
let mut set = set_prev.clone();
set.insert(i);
let reference = np.entry(set).or_insert(0);
*reference = max(*reference, a * val_prev + b);
ans = max(ans, a * val_prev + b);
}
}
dp = np;
}
println!("{}", ans);
}
計算量
事前計算では、N個の要素を定数($50$)回ソートするため$O(N \log N)$となります。
一方、後半for _ in 0..K {}のループ部分の計算量については、よくわかりません。2
一見dpの要素数が最大で$K^K$ぐらいになりそうです。
が、実際に試してみると、前計算にかかる時間のほうが支配的で、この部分の寄与はほとんど影響がないようです。
おそらくですが、ほとんどの場合、数回のループで$x>50$となってしまうため、$dp$のキーがほとんど同じになってしまうのではないかと。