[编程题]病毒检测（牛客网）

题目链接：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 
#include 
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;
}