読者です 読者をやめる 読者になる 読者になる

ゆらのふなびと

競プロ, Python, C++

yukicoder No.346 チワワ数え上げ問題(★2) 解説

yukicoderで出題していただいた問題の解説です。

問題

No.346 チワワ数え上げ問題 - yukicoder

与えられる文字列  S の部分列(連続でなくともかまわない)のうち、"cww" となるものの総数を求めてください。

制約

 1 \leq S \leq 100000

解説

計算量を落とす過程に沿って説明していきます。

 O(N^3) 解法

この問題は以下のように言い換えることができます。

文字列 Sから3文字取り出しもとの順を保って並べたとき、"cww" となるものは何通りか

そこで文字列 Sから3文字取り出す方法をすべて試し、そのそれぞれについて "cww" となるかを判定する、という方法が考えられます。ですがこれだと、計算量が O(N^3)となり間に合いません。

 O(N^2) 解法

アルゴリズムを効率化するため、「最小チワワ問題」と同じように、先頭の 'c' を固定して考えます。すると、この 'c' を使ってできる "cww" の数は、'c' よりも右側にある 'w' から2個選ぶ場合の数と一致します。

よって各 'c' について自身よりも右側にある 'w' の個数 mをカウントし、 {}_m C_{2}を加算していけばいいことになります。

ただしこれを文字列の左からやっていると O(N^2)となりまだ間に合いません。

 O(N) 解法

そこで、最小チワワ問題の O(N) 解法と同様に、文字列を右から走査します。今回は、走査しながら行う処理は以下のようになります。

  1. 今の文字が 'w' → 'w' の個数 mをインクリメント
  2. 今の文字が 'c' → sumに {}_m C_{2}を加算
  • 解答例: C++
#include <iostream>
#include <string>
using namespace std;

#define FOR(i,a,b) for (int i=(a);i<(b);i++)
#define FORR(i,a,b) for (int i=(b)-1;i>=(a);i--)
#define REP(i,n) for (int i=0;i<(n);i++)
#define RREP(i,n) for (int i=(n)-1;i>=0;i--)
#define pb push_back
#define ALL(a) (a).begin(),(a).end()

typedef long long ll;

int main() {
    string S;
    cin >> S;
    ll sm = 0;
    ll w = 0;
    RREP(i, S.length()) {
        if (S[i] == 'w') {
            w++;
        }
        if (S[i] == 'c') {
            sm += w * (w - 1) / 2;
        }
    }
    cout << sm << endl;
    return 0;
}

以下のコードでは、先にwの個数をカウントしてから {}_m C_{2}を足し上げています(文字列を右から2回走査している)

# -*- coding: utf-8 -*-

def nC2(n):
    return n * (n - 1) // 2

S = input()
N = len(S)

# w[i] = 区間[i, N]に含まれる'w'の個数
w = [0] * (N + 1)
for i in range(N)[::-1]:
    w[i] += w[i+1]
    if S[i] == 'w':
        w[i] += 1

sm = 0
for i in range(N-1):
    if S[i] == 'c':
        sm += nC2(w[i+1])
print(sm)

最小チワワ問題で O(N)解法に気づいた人には簡単だったかもしれません。