Ones and Twos (Codeforces 1896D)
Tags: Math, divide and conquer

You are given a $$$1$$$-indexed array $$$a$$$ of length $$$n$$$ where each element is $$$1$$$ or $$$2$$$.Process $$$q$$$ queries of the following two types:   "1 s": check if there exists a subarray$$$^{\dagger}$$$ of $$$a$$$ whose sum equals to $$$s$$$.  "2 i v": change $$$a_i$$$ to $$$v$$$. $$$^{\dagger}$$$ An array $$$b$$$ is a subarray of an array $$$a$$$ if $$$b$$$ can be obtained from $$$a$$$ by deletion of several (possibly, zero or all) elements from the beginning and several (possibly, zero or all) elements from the end. In particular, an array is a subarray of itself.

Input: Each test contains multiple test cases. The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of test cases. The description of the test cases follows.The first line of each test case contains two integers $$$n$$$ and $$$q$$$ ($$$1\le n,q\le 10^5$$$) — the length of array $$$a$$$ and the number of queries.The second line of each test case contains $$$n$$$ integers $$$a_1,a_2,\ldots,a_n$$$ ($$$a_i$$$ is $$$1$$$ or $$$2$$$) — the elements of array $$$a$$$.Each of the following $$$q$$$ lines of each test case contains some integers. The first integer $$$\mathrm{op}$$$ is either $$$1$$$ or $$$2$$$.   If $$$\mathrm{op}$$$ is $$$1$$$, it is followed by one integer $$$s$$$ ($$$1 \leq s \leq 2n$$$).  If $$$\mathrm{op}$$$ is $$$2$$$, it is followed by two integers $$$i$$$ and $$$v$$$ ($$$1 \leq i \leq n$$$, $$$v$$$ is $$$1$$$ or $$$2$$$). It is guaranteed that the sum of $$$n$$$ and the sum of $$$q$$$ over all test cases both do not exceed $$$10^5$$$.

Output: For each query with $$$\mathrm{op}=1$$$, output "YES" in one line if there exists a subarray of $$$a$$$ whose sum is equals to $$$s$$$, otherwise output "NO".You can output the answer in any case (upper or lower). For example, the strings "yEs", "yes", "Yes", and "YES" will be recognized as positive responses.

Note: Consider the first example:   The answer for the first query is "YES" because $$$a_1+a_2+a_3=2+1+2=5$$$.  The answer for the second query is "YES" because $$$a_1+a_2+a_3+a_4=2+1+2+1=6$$$.  The answer for the third query is "NO" because we cannot find any subarray of $$$a$$$ whose sum is $$$7$$$.  After the fourth query, the array $$$a$$$ becomes $$$[2,1,2,2,2]$$$.  The answer for the fifth query is "YES" because $$$a_2+a_3+a_4+a_5=1+2+2+2=7$$$. 

Sample Inputs:
2
5 5
2 1 2 1 2
1 5
1 6
1 7
2 4 2
1 7
3 2
2 2 2
1 6
1 5

Sample Outputs:
YES
YES
NO
YES
YES
NO

Solution (Python):
""""""
import io, os
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
import heapq
number_tests = int(input())
for _ in range(number_tests):
    min_heap = []
    max_heap = []
    value_list = set()
    (n, q) = map(int, input().split())
    a = [0] + list(map(int, input().split()))
    sum_a = sum(a)
    for i in range(1, len(a)):
        if a[i] == 1:
            min_heap.append(i)
            max_heap.append(-i)
            value_list.add(i)
    heapq.heapify(min_heap)
    heapq.heapify(max_heap)
    for _ in range(q):
        query = list(map(int, input().split()))
        if query[0] == 1:
            s = query[1]
            min_1 = None
            max_1 = None
            while len(min_heap) > 0 and min_heap[0] not in value_list:
                heapq.heappop(min_heap)
            if len(min_heap) > 0:
                min_1 = min_heap[0]
            while len(max_heap) > 0 and -max_heap[0] not in value_list:
                heapq.heappop(max_heap)
            if len(max_heap) > 0:
                max_1 = -max_heap[0]
            if min_1 == None:
                min_2 = len(a)
            else:
                min_2 = min(min_1, len(a) - max_1)
            num_2 = min_2 - 1
            k = sum_a - num_2 * 2
            if s <= k:
                print('YES')
            else:
                s -= k
                if s % 2 == 0 and s <= 2 * num_2:
                    print('YES')
                else:
                    print('NO')
        elif query[0] == 2:
            (idx, value) = query[1:]
            sum_a += value - a[idx]
            if a[idx] != value:
                a[idx] = value
                if value == 1:
                    heapq.heappush(min_heap, idx)
                    heapq.heappush(max_heap, -idx)
                    value_list.add(idx)
                else:
                    value_list.remove(idx)

Solution (Python):
""""""
import io
import os
from collections import defaultdict
import heapq

def prepro(seqs):
    global counter, maxValue, oneIndices
    counter = defaultdict(int)
    oneIndices = set()
    for (idx, seq) in enumerate(seqs):
        counter[seq] += 1
        if seq == 1:
            heapq.heappush(minIdx, idx)
            heapq.heappush(maxIdx, -idx)
            oneIndices.add(idx)
    maxValue = counter[2] * 2 + counter[1] * 1

def Query():
    minOne = None if minIdx == [] else minIdx[0]
    maxOne = None if maxIdx == [] else -maxIdx[0]
    if counter[1] == 0:
        boolean = query[0] <= maxValue and query[0] % 2 == 0
        boolean = 'YES' if boolean else 'NO'
        print(boolean)
    else:
        (left, right) = (minOne, n - (maxOne + 1))
        minDis = min(left, right)
        stick = maxValue - minDis * 2
        boolean = (query[0] <= stick or (query[0] - stick) % 2 == 0) and query[0] <= maxValue
        boolean = 'YES' if boolean else 'NO'
        print(boolean)

def update(idx, v):
    global counter, maxValue
    old = sequence[idx - 1]
    sequence[idx - 1] = v
    if old == v:
        pass
    else:
        counter[old] -= 1
        counter[v] += 1
        if v == 2:
            if idx - 1 in oneIndices:
                oneIndices.remove(idx - 1)
            if maxIdx != [] and idx - 1 == -maxIdx[0]:
                heapq.heappop(maxIdx)
                while maxIdx != [] and -maxIdx[0] not in oneIndices:
                    heapq.heappop(maxIdx)
            if minIdx != [] and idx - 1 == minIdx[0]:
                heapq.heappop(minIdx)
                while minIdx != [] and minIdx[0] not in oneIndices:
                    heapq.heappop(minIdx)
        else:
            oneIndices.add(idx - 1)
            heapq.heappush(minIdx, idx - 1)
            heapq.heappush(maxIdx, -(idx - 1))
        maxValue = counter[2] * 2 + counter[1] * 1
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
number_testcase = int(input().decode())
c = 0
for _ in range(number_testcase):
    (n, toQueries) = list(map(int, input().decode().split()))
    sequence = list(map(int, input().decode().split()))
    (minIdx, maxIdx) = ([], [])
    prepro(sequence)
    for _ in range(toQueries):
        (type, *query) = list(map(int, input().decode().split()))
        if type == 1:
            c += 1
            Query()
        else:
            (idx, v) = query
            update(idx, v)

Solution (Python):
import sys
from heapq import *
for _ in range(int(input())):
    (n, q) = map(int, input().split())
    l = list(map(int, input().split()))
    l1 = [a for (a, b) in enumerate(l) if b == 1]
    l1i = [-a for a in l1]
    heapify(l1)
    heapify(l1i)
    su = sum(l)
    if l1:
        fi = heappop(l1)
        la = -heappop(l1i)
    else:
        (la, fi) = (-1, n)
    for k in range(q):
        li = list(map(int, input().split()))
        if li[0] == 1:
            s = li[1]
            if s > su:
                print('NO')
            elif s % 2 == su % 2:
                print('YES')
            else:
                print('YES' if s < su - 2 * min(fi, n - 1 - la) else 'NO')
        else:
            (_, i, v) = li
            i -= 1
            if v == 1:
                if i < fi:
                    if fi != n:
                        heappush(l1, fi)
                    fi = i
                else:
                    heappush(l1, i)
                if i > la:
                    if la != -1:
                        heappush(l1i, -la)
                    la = i
                else:
                    heappush(l1i, -i)
            su += v - l[i]
            l[i] = v
            if v == 2 and i == fi:
                while l1:
                    u = heappop(l1)
                    if l[u] == 1:
                        fi = u
                        break
                else:
                    fi = n
            if v == 2 and i == la:
                while l1i:
                    u = -heappop(l1i)
                    if l[u] == 1:
                        la = u
                        break
                else:
                    la = -1