SRM 673 Div1 Easy BearCavalry - たこし++の備忘録

問題概要

兵士と馬の戦闘力がそれぞれwarriors、horsesという配列で与えられる。兵士と馬の数は一致する。
兵士iと馬jを組み合わせた時の戦闘力はwarriors[i]*horses[j]で表される。
兵士と馬を適当に組み合わせた時に、warriors[0](以下、リーダーとする)の戦闘力がどのユニットよりも大きくなるような組み合わせ方は全部で何通りあるかのmod $10^{9}+7$ を求めよ。

$兵士の数 \le 50$ $1 \le 戦闘力 \le 1000$

[サンプル]
warriors = {10,2,10}
horses = {100,150,200}
[答え]
3
組み合わせとしては
(10,200),(2,150),(10,100)
(10,200),(2,100),(10,150)
(10,150),(2,200),(10,100)
が考えられる。
それ以外の組み合わせは例えば(10,150)(2,100)(10,200)を考えると、リーダーの戦闘力は1500,warriors[2]の戦闘力は2000となってしまい題意を満たさないのでNG

解法

他の人のコードを見てると自分のやり方でやっている人が見当たらないのであんまり参考にならないかもしれません。

まず、リーダーに割り当てる馬を適当に一つ決め、残りの兵士と馬でその戦闘力未満となるような割り当て方がいくつあるかを数える。

馬の戦闘力をソートしておき、各兵士[i]が乗ることのできる馬の番号の最大値を配列Width[i]に保存し、ソートをすると、「各兵士[i]が座ることのできる範囲が[1,Width[i]]で与えられ、i $<$ jの時Width[i] $\le$ Width[j]が保証されている時の兵士の座り方の組み合わせ数」を求める問題になる。

これは、dp[i+1][j][k] := i番目までの兵士を見た時に一番右に座っている兵士の座席をj、j未満でまだ兵士が座っていない席の数をkとした時の座り方の数
としたdpを更新していけば解ける。

dpの更新でO( $N^{4}$ )で、最初のリーダーに馬を割り当てる割り当て方はN通りあるので、全体の計算量はO( $N^{5}$ )
計算量はかなり多めだけど、サーバーのジャッジが早いので通ります。

コード

#include <cstdio>
#include <cmath>
#include <cstring>
#include <ctime>
#include <iostream>
#include <algorithm>
#include <set>
#include <vector>
#include <sstream>
#include <typeinfo>
#include <fstream>

using namespace std;

#define mod 1000000007
#define MAX_N 55

typedef long long LL;

//dp[i+1][j][k] := i番目の兵士まで見た時、
//既に割り当てた一番右にある馬の番号をjとし、
//j未満でまだ余っている馬の数をkとしたときの
//組み合わせの数
LL dp[MAX_N][MAX_N][MAX_N];

class BearCavalry {
public:

  //各兵士が乗ることのできる馬の数
  int Width[MAX_N];

  LL Solve(int X, vector<int> warriors, vector<int> horses){

    //馬をソートしておく
    sort(horses.begin(), horses.end());

    int N = warriors.size();

    //各兵士が乗ることのできる馬の数を求めておく
    for(int i = 0; i < N; i++){
      bool flag = false;
      for(int j = N-1; j >= 0; j--){
	if(warriors[i]*horses[j] < X){
	  Width[i] = j+1;
	  flag = true;
	  break;
	}
      }
      if(!flag)
	return 0;
    }

    //各兵士が乗ることのできる馬の数でソートする
    sort(Width, Width+N);

    memset(dp, 0, sizeof(dp));

    dp[0][0][0] = 1;

    for(int i = 0; i < N; i++){
      for(int j = 0; j <= N; j++){
	for(int k = 0; k <= N; k++){
	  
	  if(dp[i][j][k] == 0)
	    continue;

	  //今まで割り当てた一番右の馬を更新できるとき
	  if(j+1 <= Width[i]){
	    for(int l = j+1; l <= Width[i]; l++){
	      dp[i+1][l][k + (l - j - 1)] = (dp[i+1][l][k + (l - j - 1)] +
					     dp[i][j][k]) % mod;
	    }
	  }

	  //今まで割り当てた一番右の馬より左に、
	  //まだ割り当てていなう馬にその兵士を割り当てるとき
	  if(k > 0){
	    dp[i+1][j][k-1] = (dp[i+1][j][k-1] + 
			       (dp[i][j][k]*k)%mod) % mod;
	  }

	}
      }
    }

    return dp[N][N][0] % mod;

  }

  int countAssignments(vector<int> warriors, vector<int> horses) {

    //BraveBeartをwarriorsから取り除く
    int brave = warriors[0];
    warriors.erase(warriors.begin());

    int ret = 0;

    //BraveBeartに割り当てる馬を全探索する
    for(int i = 0; i < horses.size(); i++){
      int tmp = horses[0];
      //BraveBeartに割り当てる馬は取り除く
      horses.erase(horses.begin());
      ret = ((LL)ret + Solve(tmp*brave, warriors, horses)) % mod;
      //馬を戻しとく
      horses.push_back(tmp);
    }
    
    return ret;
  }
};