CFR 519 D. A and B and Interesting Substrings ( Ad hoc )

Problem - 519D - Codeforces

題意：
給 26 個英文字母各自的花費，接著一個字串。求有多少子字串，滿足頭尾字母相同，且扣去頭尾字母後的花費總和為 0。

數據範圍：
花費 -1e5 ≤ cost[ i ] ≤ 1e5
字串長度 1 ≤ | seq | ≤ 1e5

解法：
先預處理前綴花費和，並做記錄每個字母的出現位置。接著對每個字母分別計算，由出現位置前面至後面更新答案，每次從 map 取出有多少前面的前綴和等於當前這個地方前一個的前綴和，因為這等價去頭尾字母厚的花費總和為 0 的情況，更新完答案後更新 map。

時間 / 空間複雜度：
O( | seq | lg | seq | )

vi cost;
string seq;

void init(){
    cost = vi( 256 );
    for( int i = 'a'; i <= 'z'; ++i )
        cin >> cost[ i ];
    cin >> seq;
}

vl pscost;
vvi ch_pos;

void preprocess(){
    pscost = vl( seq.size() + 1 );
    for( int i = 0; i < seq.size(); ++i )
        pscost[ i + 1 ] = pscost[ i ] + cost[ seq[ i ] ];
    ch_pos = vvi( 256 );
    for( int i = 0; i < seq.size(); ++i )
        ch_pos[ seq[ i ] ].emplace_back( i );
}

void solve(){
    ll ans = 0;
    for( int i = 'a'; i <= 'z'; ++i ){
        map< ll, int > cnt;
        for( int j = 0; j < ch_pos[ i ].size(); ++j )
            ans += cnt[ pscost[ ch_pos[ i ][ j ] ] ],
            ++cnt[ pscost[ ch_pos[ i ][ j ] + 1 ] ];
    }
    cout << ans << endl;
}