Atcoder解いてみた AtCoder Beginner Contest E - Active Infants

AtCoder Beginner Contest E - Active Infants

問題はこちら

回答

素直な全探索では、TLE
O(N!)となり、全く収まらない

Aiを大きい順にソートし、たその時点で得点が大きい側へ移動させる貪欲法の場合はWAとなる
例えば、
Aiが1,2,3,4,5,6のように与えられた場合、
6,5に関しては貪欲に左隅において大丈夫だが、4に関しては、貪欲に右隅に置くのではなく、左隅に置き、その後1,2を右隅に置いた方がいい

Aiの大きい方から配置していくのはいいのだが、貪欲に位置を決めるのではなく、左・右への移動も考慮する必要

N<=2000なため、O(N^2)あたりまでなら計算量として全く問題ない
加えて、左隅 or 右隅に置くことを考慮しつつ(どの座標かはいらない)、価値の最大値を管理できればよいため、以下の配列を用いてDPを行うことを考える

dp[l][r] : 左隅にl個、右隅にr個置いた時の価値の最大値
l,rそれぞれでループを回し、A[i]を最初から見ていき、A[i]を左隅に置いた時の価値、右隅に置いた時の価値を更新していけばいい

ちなみに、貪欲法の反例探しだが、今これを選ぶことで、後の人が困るケースを作れないか？と考えるといい
今回の場合は、この方針にのっとり、Aiの差分、距離の差分が小さい状況を考えることで反例を見つけることができる
完全に反例を見つけられなくても、貪欲法が怪しいと気づくために、Aiが1,2,3,4,9,5,8の時、9を左端にやるよりも、8を左端にやった方が得なため、貪欲に値を左端 or 右端のどちらに置けばいいのか決めることはできないと観察する方向でもいい
すると、あり得る状態遷移を考慮して、各状態の最適な値を更新する方向を疑い、DPに繋げられるかもしれない

#include <iostream>
#include <string>
#include <map>
#include <unordered_map>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <limits.h>
#include <bitset>
#include <list>
#include <set>
#include <numeric>
#include <tuple>

int N;

int main()
{
	std::cin.tie(0);
	std::ios::sync_with_stdio(false);

	std::cin >> N;
	std::vector<std::pair<long long, long long>> Ai;

	for (int i = 0; i < N; i++) {
		std::pair<long long, long long> p;
		std::cin >> p.first;
		p.second = i;

		Ai.push_back(p);
	}

	std::sort(Ai.begin(), Ai.end(), std::greater<std::pair<long long, long long>>()); // Aiの大きい順に処理


	// dp[l][r] : 左にl個、右にr個置いたときの価値の最大値
	std::vector<std::vector<long long>> dp(N + 1, std::vector<long long>(N + 1, LLONG_MIN));

	dp[0][0] = 0; // 左右どちらにも配置していないため、0

	for (int l = 0; l <= N; l++) {
		for (int r = 0; r <= N; r++) {

			// まだスコアが定められていない(状態遷移できていない)場合、スキップ
			// これがなくてもACすることがあるが、後ろで、dp[l][r]使って計算しているため、明示的に、到達していない状態の場合はスキップする処理入れておいたほうがいい
			if (dp[l][r] == LLONG_MIN) {
				continue;
			}

			int aIdx = l + r;

			// 左右で計N個より多く配置できないことを考慮
			// ただ、Aiのインデックスに合わせる
			if (aIdx >= N) {
				continue;
			}

			// 左の配置位置
			int left = l;

			// 右の配置位置
			int right = N - r - 1;

			long long leftScore = Ai[aIdx].first * std::abs(left - Ai[aIdx].second);
			long long rightScore = Ai[aIdx].first * std::abs(right - Ai[aIdx].second);

			// 左に配置する場合
			dp[l + 1][r] = std::max(dp[l][r] + leftScore, dp[l+1][r]);

			// 右に配置する場合
			dp[l][r + 1] = std::max(dp[l][r] + rightScore, dp[l][r+1]);
		}
	}

	long long ans = 0;

	// dp内のl+r=Nのものから選ぶ
	for (int l = 0; l <= N; l++) {
		for (int r = 0; r <= N; r++) {
			if (l + r != N) {
				continue;
			}
			ans = std::max(ans, dp[l][r]);
		}
	}

	std::cout << ans << std::endl;
	return 0;
}