Submission #50726374

---

Source Code Expand

#include <bits/stdc++.h>
#include <atcoder/modint>

namespace {
    using ModInt [[maybe_unused]] = atcoder::modint998244353;
    using Num [[maybe_unused]] = long long int;
    using Vec [[maybe_unused]] = std::vector<Num>;
    using Set [[maybe_unused]] = std::set<Num>;
    using Mset [[maybe_unused]] = std::multiset<Num>;
    using Edges [[maybe_unused]] = std::vector<std::vector<Num>>;

    template<typename T>
    using Q [[maybe_unused]] = std::queue<T>;

    template<typename T>
    using PQ [[maybe_unused]] = std::priority_queue<T, std::vector<T>, std::greater<T>>;
}

void solve(std::istream& is, std::ostream& os) {
    Num loop {0};
    is >> loop;

    while(loop-- > 0) {
        Num n, k;
        std::string s;
        is >> n >> k >> s;

        char prev = '0';
        std::vector<std::pair<Num, Num>> runs;
        Num n_ones {0};

        for(size_t i{0}; i<s.size(); ++i) {
            const auto c = s.at(i);
            if (c != '0') {
                if (prev == '0') {
                    runs.push_back(std::make_pair(i, 1));
                } else {
                    runs.back().second += 1;
                }
            }

            n_ones += (c == '1');
            prev = c;
        }

        std::vector<size_t> ids;
        for(size_t i{0}; i<runs.size(); ++i) {
            const auto& [left, len] = runs.at(i);
            Num n1 {0};
            for(Num j{left}; j<(left+len); ++j) {
                n1 += (s.at(j) == '1');
            }
            if (n1 > 0) {
                ids.push_back(i);
            }
        }

        if (ids.size() > 1) {
            os << "No\n";
            continue;
        }

        Num found {0};
        for(size_t id{0}; id<runs.size(); ++id) {
            const auto& [left, len] = runs.at(id);
            if (len < k) {
                continue;
            }

            if (!ids.empty() && (ids.at(0) != id)) {
                continue;
            }

            Vec cumsum {0};
            for(Num i{left}; i<(left+len); ++i) {
                Num b = 0 + (s.at(i) == '1');
                cumsum.push_back(cumsum.back() + b);
            }

            for(Num i{0}; i<=(len-k); ++i) {
                const Num left = cumsum.at(i);
                const Num right = cumsum.back() - cumsum.at(i+k);
                if ((left == 0) && (right == 0)) {
                    ++found;
                }
            }
        }

        if (found == 1) {
            os << "Yes\n";
        } else {
            os << "No\n";
        }
    }
}

int main(void) {
    solve(std::cin, std::cout);
    return 0;
}

Submission Info

Judge Result