Lilliput Steps

小さな一歩から着実に. 数学やプログラミングのことを書きます.

APIO 2010 - commando

問題文 : 奇襲部隊


解法

まず, 少し考察をすると, $f(n) = \displaystyle \sum_{i = 1}^{n}x_n$ として

$dp[i] = \max [ \underline{dp[j] + a \{ f(i) - f(j) \} ^2 + b \{f(i) - f(j) \} + c},\ a f(i)^2 + bf(i)+c ] (1 \leqq j \leqq i - 1)$

というDP の漸化式が立つ事が分かる. これをそのまま計算すると$O(N^2)$ で50 点が得られる.

ここで, dp の漸化式の前半の式(下線部)を変形すると

$dp[j] + a \{ f(i) - f(j) \} ^2 + b \{f(i) - f(j) \} + c = af(i)^2+bf(i)+c-2af(i)f(j)+af(j)^2-bf(j)+dp[j]$

となる. ここで$A_j = -2af(j)$, $x=f(i)$, $B_j=af(j)^2-bf(j)+dp[j]$ とすると, それぞれの変数は$i,\ j$ について独立で, $dp[j]$ は一次式群$A_jx+B_j$に$af(i)^2+bf(i)+c$ を加算したものの中の最大値となる.

これはconvex hull trick という考え方を使うことで$O(N)$ で計算できる.

近いうちにconvex hull trick についての記事かきます.

コード :

#include <cstdio>
#include <vector>

using namespace std;

typedef long long lint;

vector<lint> A, B;

bool check(int i, int j, int k)
{
	return ((A[i] - A[k]) * (B[j] - B[i]) < (A[i] - A[j]) * (B[k] - B[i]));
}

void add(lint a, lint b)
{
	A.push_back(a); B.push_back(b);
	
	while (A.size() >= 3 && check(A.size() - 1, A.size() - 2, A.size() - 3)){
		A.erase(A.end() - 2); B.erase(B.end() - 2);
	}
}

lint getMax(lint x)
{
	static int idx = 0;
	if (idx > A.size()) idx = A.size() - 1;
	
	while (idx < A.size() - 1 && A[idx] * x + B[idx] < A[idx + 1] * x + B[idx + 1])
		idx++;
	
	return (A[idx] * x + B[idx]);
}

lint dp[1000001];
lint sum[1000001];

int main()
{
	int n;
	lint a, b, c;
	scanf("%d", &n);
	scanf("%lld %lld %lld", &a, &b, &c);
	
	for (int i = 0; i < n; i++){
		lint x;
		scanf("%lld", &x);
		sum[i + 1] = sum[i] + x;
	}
	
	dp[0] = 0; dp[1] = a * sum[1] * sum[1] + b * sum[1] + c;
	add(-2 * a * sum[1], dp[1] + a * sum[1] * sum[1] - b * sum[1]);
	for (int i = 2; i <= n; i++){
		lint cons = a * sum[i] * sum[i] + b * sum[i] + c;
		dp[i] = max(cons, cons + getMax(sum[i]));
		add(-2 * a * sum[i], dp[i] + a * sum[i] * sum[i] - b * sum[i]);
	}
	
	printf("%lld\n", dp[n]);
}