裏紙

ほぼ競プロ、たまに日記

CF 733 E - Sleep in Class

問題

Problem - E - Codeforces

問題概要

n段の階段がある。各段は、下から順に番号が1からnまで付けられており、その各段にはupdownかのどちらかが書かれており、その段にいるときに進む方向を示している。そして、そこから移動した瞬間に、その段の示す方向は反転する。

さて、次のような動作を考える:はじめに、あなたはi段目の位置にいて、1秒に1回の移動を行う。そして、1段目にいるときにdownするかn段目にいるときにupすると階段から落ちる。落ちるのにも1秒かかるとしたとき、落ちるまでにかかる時間を各i(1 \le i \le n)について求めよ。いつまでも落ちることのない場合は、-1と答えよ。

  •  1 \le n \le 10^6

イデア

UとDの区間が交互に連続していると考えるとその区間はどんどんマージされていくので1つの立ち位置だけなら、シミュレーションによりO(n)で解けるように感じるが、今回はそれを全スタート位置についてやらなければいけないのでそれは厳しい。

あるスタート地点に立った時のことを考えてみると、そこから何回折り返しが発生するかというのはスタート位置よりも上にあるDと下にあるUの個数に依存してくることに気づく。

ということで、先頭からのUの個数を累積和として配列kuに保存し、末尾からのDの個数を累積和としてkdに記録することにする(この累積和に関して、自分自身のマスは考慮しないことにして、それぞれ自分より小さい範囲、自分より大きい範囲の位置を考える)。

次に、次のような配列tltrを考えてみる:

  • tl_i: s_iが常にDで不変であると仮定したときに位置iからスタートして脱出までにかかる時間
  • tr_i: s_iが常にUで不変であると仮定したときに位置iからスタートして脱出までにかかる時間

tlの方に注目してみると、まずtl_1 = 1は明らか。そして、ii+1の関係を考えてみる。すると、tl_{i+1}tl_iの状況に比べてUに遭遇したときに戻ってくるのに距離が1だけ伸びていることになる。これを行き帰りで考慮しなければならないので1つのUにつき時間が2だけ加算される。それと、はじめの一歩を考慮して tl_{i+1} = tl_i + 2 * ku_{i+1} +1という漸化式が立てられる。trの方も同様にして、tr_n = 1で初期化し、 tr_{i-1} = tr_i + 2 * kd_{i-1} + 1となる。

さて、階段の一番上の段から上に登って落ちるのか、一番下の段から下に降りて落ちるのかは、ku_ikd_iの大小関係から分かる。もしこの2つの値が等しい時には、スタート地点の方向を見れば、その方向に一致する。この関係性から落ちる方向は決まっているが、上からと下から、同時に落下すると仮定して、その合計の時間から実際には脱出出来ない方の時間の一部を引くという形でスタート位置がi段目のときにかかる時間を求めていくことにする。このような方針にすることで、上で不自然な感じに設定したtltrが活かせる(iを基準位置に区切ってかかる時間を考えることが出来るようになった)。

出口の方向にもよるが、tl_i + tr_iからtl_jまたはtr_jを引くことになる(jは出口とは反対方向で最終的に到達する位置)。また、ちょうど位置jにいるときに反対で脱出されるわけではないので、数合わせをする必要もあるが、まずはjについて考えてみる。

下から脱出するとすると、kd_i - ku_i - f \ge 0になる(fs_iUなら1、Dなら0)。反対側で、最終的に到達する位置jは、この折り返しの回数を考慮すると、配列posdDの現れる位置を上から順に保存しておくことにすれば、j = posd [ kd_i - ku_i - f ] というということになる。tr_jを引けばよいわけだが、それだけでは足りない。余計にijの間を何回か往復していることになっているのでその分も引いて上げる必要がある。そして往復分を考えれば、ここで引くべきなのは (2(kd_i - ku_i - f) + 1) * (j-i)となる。これは、のこってる分だけ往復するのが第一項で、第二項が最後は結局反対側から出ていくので1回だけしか通過しないところがあるぶんを引いている。

上から脱出する時には、同様にposdの代わりにposu(Uの現れる位置を下から順に保存)を使って上と同じ考察をすればよい。

実装(C++)

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define each(itr,c) for(__typeof(c.begin()) itr=c.begin(); itr!=c.end(); ++itr)
#define all(x) (x).begin(),(x).end()
#define pb push_back
#define fi first
#define se second

const int N=1000000;

int n;
char s[N+1];

ll ku[N]={},kd[N]={};
ll tl[N],tr[N];

int posd[N],posu[N];
int PD=0,PU=0;

int main()
{
    scanf(" %d", &n);
    scanf(" %s", s);

    for(int i=1; i<n; ++i) ku[i] = ku[i-1]+(s[i-1]=='U');
    for(int i=n-2; i>=0; --i) kd[i] = kd[i+1]+(s[i+1]=='D');

    tl[0]=1;
    for(int i=1; i<n; ++i) tl[i] = tl[i-1] + 2*ku[i] + 1;
    tr[n-1]=1;
    for(int i=n-2; i>=0; --i) tr[i] = tr[i+1] + 2*kd[i] + 1;

    rep(i,n)if(s[i]=='U') posu[PU++] = i;
    for(int i=n-1; i>=0; --i)if(s[i]=='D') posd[PD++] = i;

    rep(i,n)
    {
        int f = (s[i]=='U');
        ll ans;
        if(kd[i]-ku[i]-f >= 0)
        {
            ans = tl[i]+tr[i];
            int j = posd[kd[i]-ku[i]-f];

            ans -= tr[j];
            ans -= (2LL*(kd[i]-ku[i]-f)+1)*(j-i);
        }
        else
        {
            f = (s[i]=='D');
            ans = tl[i]+tr[i];
            int j = posu[ku[i]-kd[i]-f];

            ans -= tl[j];
            ans -= (2LL*(ku[i]-kd[i]-f)+1)*(i-j);
        }

        if(i) printf(" ");
        printf("%lld", ans);
    }
    printf("\n");
    return 0;
}