[編程題]病毒檢測（牛客網）

題目鏈接：https://www.nowcoder.com/questionTerminal/6f0d16fc06274f44af8913d182668037
來源：牛客網

小明最近在做病毒自動檢測，他發現，在某些library 的代碼段的二進制表示中，如果包含子串並且恰好有k個1，就有可能有潛在的病毒。library的二進制表示可能很大，並且子串可能很多，人工分析不可能，於是他想寫個程序來先算算到底有多少個子串滿足條件。如果子串內容相同，但是開始或者結束位置不一樣，則被認爲是不同的子串。
注：子串一定是連續的。例如"010"有6個子串，分別是 "0, “1”, “0”, “01”, “10”, “010”
輸入描述:
第一行是一個整數k，表示子串中有k個1就有可能是病毒。其中 0 <= k <= 1 000 000

第二行是一個字符串，就是library的代碼部分的二進制表示。字符串長度 <= 1 000 000。並且字符串中只包含"0"或"1".
輸出描述:
輸出一個整數，所有滿足只包含k個1的子串的個數。
Sample
輸入
1
1010
輸出
6
說明
滿足條件的子串有：“1”, “1”, “10”, “01”, “10”, “010”.
初步想法：暴力求解，遍歷每個子串，再求每個子串包含1的個數，很容易想到，這需要三層循環，時間必定超出，只有50%AC。
改進思路：其實換種想法，捨去0不看，我們要做的就是找到連續個k的1區間，然後再通過這個區間前後連續的0的個數來判斷，當前區間可以組成多少符合要求的子串。
故我們可以設置兩個vector，一個v1用來存儲1的位置，另一個v2用來存儲某個1前面的0的數量（最後一個值是末尾0的數量，如果給定字符串末尾是1，則值爲0）。根據前面的想法，我們利用一個長度爲k的窗口在v1上滑動，再根據v2得到某個1前面或後面的0的個數（一段滿足條件的1的個數能構成的子串數量由其左右兩邊的0的數量決定，比如左邊3個0，右邊4個0，那麼自由組合一下就是(3+1)*(4+1)=20個，思路參考於牛客網評論區）。

#include <iostream>
#include <vector>
using namespace std;
int main(){
    int k;
    long long num = 0;   //有一個測試用例全是0，結果會超出
    string s;
    char c = ' ';
    vector<int> loc,zeros;  //loc記錄所有1個位置，zeros是每個1前面0的數量
    cin>>k;
    cin>>s;
    long long count = 0;
    for(int i = 0;i < s.length();++i){
        c = s[i];
        if(c == '1'){
            loc.push_back(i);
            zeros.push_back(count);
            count = 0;  //記錄完1前面的0個數後，清零
        } else{
            count++;
        }
    }
    if( c == '1') zeros.push_back(0);   //記錄最後一個1後的剩餘的0的個數
    else    zeros.push_back(count);
    if( k == 0){
        long long val = 0;
        for(int i = 0;i < s.length();++i){
            c = s[i];
            if(c == '0') val++;
            else{
                num += (val * (val + 1) / 2);
                val = 0;
            }
        }
        if(c == '0'){
            num += (val * (val + 1) / 2);
        }
    } else if (k > loc.size()) num = 0;
    else{
        for (int i = 0; i + k - 1 < loc.size(); ++i) {
            num += ((zeros[i] + 1) * (zeros[i+k] + 1));
        }
    }
    cout<<num;
    return 0;
}