累積和を用いて高速化をする

Question

質問内容

AtCoderの過去問を解いていて，答えが一致するが実行時間が制限以内でない判定を受けてしまったので解説を見てみたのですが，
西から i 番目の人がリーダーになった際、向く方向を変える必要のある人は、西から j(j < i) 番
目に並んでいて西を向いている人と、西から j(j > i) 番目に並んでいて東を向いている人です。
ここで、「ある範囲に存在する西 (東) を向いている人の人数」は、累積和を使うと、前処理 O(N)
時間、1 クエリあたり O(N) 時間で求められます。
よって、まず最初に O(N) 時間で累積和を求めておくと、各人について、その人がリーダーに
なった際に向く方向を変える必要のある人の人数を求めることが O(1) 時間ででき、合計 O(N) 時
間で全員分求まります。あとはその中で最小値を求めればよいです。

とあって累積和を使うとうまくいくそうなのですが，どのような実装の仕方をすればいいのか教えていただきたいです．自分で考えたものを以下に載せますが，これが自分的には累積和になっているのではないかと思ってしまっているので，なにが累積和とは違うのか教えていただきたいです．

for(int i=0;i<N;i++){
        if(str[i]=='W'){
            for(int j=i+1;j<N;j++) count[j]++;
        }
        else{
            for(int j=0;j<i;j++) count[j]++;
        }
    }

問題

コード全体

#include<stdio.h>
#include<stdlib.h>
int main(void){
    int N;
    char *str;
    int *count;
    int min;
    scanf("%d",&N);
    min=N;
    str=(char *)malloc(sizeof(char)*N);
    count=(int *)malloc(sizeof(int)*N);
    scanf("%s",str);
    //printf("%s\n",str);

    for(int i=0;i<N;i++) count[i]=0;

    for(int i=0;i<N;i++){
        if(str[i]=='W'){
            for(int j=i+1;j<N;j++) count[j]++;
        }
        else{
            for(int j=0;j<i;j++) count[j]++;
        }
    }

    for(int i=0;i<N;i++){
        if(min>count[i]) min=count[i];
    }
    printf("%d\n",min);
    
    return 0;
}

Answer 1

なにが累積和とは違うのか

計算量が違います。解説では累積和の計算がO(N)時間で終わることになってますが、そのコードでは二重ループになっているのでO(N^2)かかります。

O(N)で計算するには、すでに計算されてるcount[i - 1]の値を利用して、

count[i] = count[i - 1] + 変化分

とするだけですが、今のコードのように一つの配列で済まそうとするとなかなかややこしいので、自分の実装力と相談してください。

---

count[i] は[0, i)に含まれる'W'の個数と(i, N)に含まれる'E'の個数の和になるはずです。つまりちょうどi番目の'W'も'E'も含まれてはいけないはずですが、count[0]の時点で0番目の'E'がカウントされてます。これが原因でしょう。

Comments

[grape_ll]

解説ありがとうございます．
教えていただいたことをもとにコードを書いて提出してみましたら，時間は圧倒的に早くなって完全にクリアしました．しかし，結果が答えと一致しない問題が半数ほど出てきてしまいました．
もしよろしければどこがよくないのか教えていただきたいです．

#include<stdio.h>
#include<stdlib.h>
int main(void){
int N;
char *str;
int *count;
int min;
scanf("%d",&N);
min=N;
str=(char *)malloc(sizeof(char)*N);
count=(int *)malloc(sizeof(int)*N);
scanf("%s",str);
//printf("%s\n",str);

for(int i=0;i<N;i++) count[i]=0;

for(int i=0;i<N;i++){
if(str[i]=='E') count[0]++;
}
for(int i=1;i<N;i++){
if(str[i-1]=='W' && str[i]=='E'){
count[i]=count[i-1];
}
else if(str[i-1]=='E' && str[i]=='E'){
count[i]=count[i-1]-1;
}
else if(str[i-1]=='W' && str[i]=='W'){
count[i]=count[i-1]+1;
}
else if(str[i-1]=='E' && str[i]=='W'){
count[i]=count[i-1];
}
}

for(int i=0;i<N;i++){
if(min>count[i]) min=count[i];
}
printf("%d\n",min);

return 0;
}

Answer 2

e.g.

←西                 東→
0 1 2 3 4 5 6 7 8 9 A B  //index
-----------------------
W E E E W E E E W W W E  //Data
-----------------------
1 1 1 1 2 2 2 2 3 4 5 5  //Wの累積(西側から)
7 7 6 5 4 4 3 2 1 1 1 1  //Eの累積(東側から)

リーダーのindexが5であるとき，

• リーダーよりも西にいて，西を向いている者の人数は，Wの累積の[4(=5-1)]の値
• リーダーよりも東にいて，東を向いている者の人数は，Eの累積の[6(=5+1)]の値