ぱーぽーの競プロ記

競技プログラミングに関することを書きます。

CODE THANKS FESTIVAL 2014 A日程 G : 通勤電車と気分

概要

電車には1~Kの番号が付けられたK個の座席が一列に並んでおり、N人がこれらの座席に座ることを考える。

人は以下の2種類の座席の選び方のどちらかによって座る座席を選ぶ。

  • 「とにかく座りたい気分」:空席の中で最も番号が小さいものを選ぶ。
  • 「余裕があれば座りたい気分」:両隣が空席であるような空席の中で最も番号が小さいものを選ぶ。隣の座席がない場合も空席とみなす。

座席の選び方は人それぞれ異なり、I番目の人はPiパーセントの確率で「とにかく座りたい気分」となる。

空席の期待値を求めよ。

N<=100
K<=200

http://code-thanks-festival-2014-a-open.contest.atcoder.jp/tasks/code_thanks_festival_14_quala_g

解法

動的計画法で以下のような式で求めることができる。

dp[i番目][最も右に座っている人より左の空席数][最も右に座っている人]:=そのときの確率

(座席の番号が小さいほうを左側としている)

「余裕があれば座りたい気分」の両隣が空席かどうかの処理は、上記式の最も右に座っている人の隣の隣の座席に座らせるようにする。

DPで確率を計算することができれば、そこから期待値を求めることができる。

余談

vector同士のswap機能があることを初めて知った。
vectorで配列を再利用しながらDPするときに使えて便利そうだ。

データの型はswap対象が同じでないといけないが、要素数は等しくなくてもよいらしい。

ソースコード

#include <bits/stdc++.h>

#define REP(i, x, n) for(int i = x; i < (int)(n); i++)
#define rep(i, n) REP(i, 0, n)
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
#define F first
#define S second
#define mp make_pair

using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> P;

int N, K;
int p[100];
double dp[101][201][201];
// dp[i番目][最も右に座っている人より左の空席数][最も右に座っている人]

int main() {
  // ios_base::sync_with_stdio(false);
  cin >> N >> K;
  rep(i, N) cin >> p[i];

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

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

        double pb = (double)p[i] / 100;
        
        // とにかく座りたい気分
        if(j - 1 >= 0) {
          dp[i + 1][j - 1][k] += dp[i][j][k] * pb;
        }
        else if(k + 1 <= K) {
          dp[i + 1][j][k + 1] += dp[i][j][k] * pb;
        }
        else {
          dp[i + 1][j][k] += dp[i][j][k] * pb;
        }

        // 余裕があれば座りたい気分
        if(k == 0) {
          dp[i + 1][j][k + 1] += dp[i][j][k] * (1 - pb);
        }
        else if(k + 2 <= K) {
          dp[i + 1][j + 1][k + 2] += dp[i][j][k] * (1 - pb);
        }
        else {
          dp[i + 1][j][k] += dp[i][j][k] * (1 - pb);
        }
      }
    }
  }
  
  double ans = 0;
  for(int j = 0; j <= K; j++) {
    for(int k = 0; k <= K; k++) {
      ans += dp[N][j][k] * (j + K - k);
    }
  }

  printf("%.10f\n", ans);
  
  return 0;
}