SRM 513 250pt: YetAnotherIncredibleMachine

問題概要

リンゴが落ちるゲーム(リンゴの個数=N(<50))で、板がM(<50)枚ある。各板の長さと各板にカバーされていないとならない座標が指定されるので、全てのリンゴが落ちるような場合の数を求める問題。

考えたこと

  • 各板は独立に考えてよい。
  • 板の長さは大したことないので、1ずつ動かす方針でやろう(O(N*M*LENGTH))。
  • これ計算量下げることもできる(O(N*M))だろうけど、植木算っぽくなって確実にバグ埋め込むのでそんなことはしない。間に合う限り最もシンプルに書くべきだ。

本番中に書いたコード

#include <vector>
using namespace std;

typedef long long int64;

const int64 MOD = 1000000009LL;

struct YetAnotherIncredibleMachine {

	int countWays(vector <int> platformMount, vector <int> platformLength, vector <int> balls) {
		int64 ret = 1;
		int N = platformLength.size();
		int M = balls.size();
		for(int i=0; i<N; i++){
			int cnt = 0;
			for(int k=-platformLength[i]; k<=0; k++){
				bool ok = true;
				for(int j=0; j<M; j++){
					if(platformMount[i] + k <= balls[j] && balls[j] <= platformMount[i] + k + platformLength[i]){
						ok = false;
					}
				}
				if(ok){
					cnt++;
				}
			}
			ret = (ret * cnt)%MOD;
		}
		return (int)ret;
	}

};

バグを減らすための工夫

途中で0が出たらその時点でreturn 0する、というのは最悪計算量のテストがやりにくくなるとか、コードの長さが無駄に長くなってしまうのでやらない。
こういう、全て独立なことが分かっているものは関数に分けて書くのが好ましいような気もする。ネストが深くなるとコードが読みづらくなるので。それを反映したコードは次のような感じになる。
変数okの更新の仕方が微妙に変わっているけど、これは好みの問題だと思う。

struct YetAnotherIncredibleMachine {

	int64 func(int mount, int length, const vector<int>& balls){
		int M = balls.size();
		int64 cnt = 0;
		for(int left = mount - length; left <= mount; left++){
			bool ok = true;
			for(int i=0; i<M; i++){
				ok &= !(left <= balls[i] && balls[i] <= left + length);
			}
			if(ok) cnt++;
		}
		return cnt;
	}

	int countWays(vector <int> platformMount, vector <int> platformLength, vector <int> balls) {
		int64 ret = 1;
		int N = platformLength.size();
		for(int i=0; i<N; i++){
			ret = (ret * func(platformMount[i], platformLength[i], balls))%MOD;
		}
		return (int)ret;
	}

};

計算量を下げる

O(N*M)解法

この問題では、よく考えると有効な区間は必ず連続していることが分かる。したがって、「最も左にどこまで動けるか」と「最も右にどこまで動けるか」を求めてやればよい。
それを反映したのが次のコードだけど、この手のは端の処理を考えるのが面倒(+1したり-1したり)というか、うっかり間違えやすい。植木算とか苦手です。

#include <algorithm>

	int64 func2(int mount, int length, const vector<int>& balls){
		int M = balls.size();
		int left = mount - length, right = mount; //どちらも左端の座標
		for(int i=0; i<M; i++){

			//leftを更新
			if(mount - length <= balls[i] && balls[i] <= mount){
				left = max(left, balls[i] + 1);
			}

			//rightを更新
			if(mount <= balls[i] && balls[i] <= mount + length){
				right = min(right, balls[i] - length - 1);
			}
		}

		return max(0, right - left + 1);
	}
O((N+M)*logM)解法

上のとアイデアは同じで、ボールをあらかじめソートして置けば最左、最右を決める座標が二分探索で求められて嬉しい。lower_boundの使い方に注意。その値以下で最大のものを探すにはどうするか、とかはlower_bound使うかupper_bound使うか降順のlower_bound使うかとかを間違えやすい。二分探索を自分で書くよりは安全だけど。あと、const_iteratorとかconst_reverse_iteratorとか初めて使った。

#include <algorithm>
#include <functional>

        //前処理としてballsをソートしておく
	int64 func3(int mount, int length, const vector<int>& balls){
		int left = mount - length, right = mount; //どちらも左端の座標

		//leftを更新
		{
			const vector<int>::const_reverse_iterator itr = lower_bound(balls.rbegin(), balls.rend(), mount, greater<int>());
			if(itr != balls.rend() && mount - length <= *itr && *itr <= mount){
				left = (*itr) + 1;
			}
		}

		//rightを更新
		{
			const vector<int>::const_iterator itr = lower_bound(balls.begin(), balls.end(), mount);
			if(itr != balls.end() && mount <= *itr && *itr <= mount + length){
				right = (*itr) - 1 - length;
			}
		}

		return max(0, right - left + 1);
	}