如何高效实现全1数组的置0操作与第k个1定位查询？

首先，先解决你二分查找方案的问题，再介绍更适合大规模场景的高效算法。

修复你的二分查找方案 #

你的核心思路是对的——维护一个存储所有1的原索引的有序列表，查询时用二分快速定位第k个元素，但类型0操作的逻辑错误导致了索引混乱：你直接用输入的b（原数组的1-based索引）作为列表下标去pop，但实际上应该先找到原索引（b-1）在列表中的位置，再删除它。

因为这个列表始终是升序的（初始是0到n-1，删除元素后依然保持有序），我们可以用bisect模块快速定位目标索引的位置：

import bisect

n = int(input())
q = int(input())
# 初始存储所有1的0-based原索引，保持升序
ones_indices = list(range(n))

for _ in range(q):
    op, val = map(int, input().split())
    if op == 0:
        # 转换为原数组的0-based索引
        target = val - 1
        # 找到target在列表中的位置
        pos = bisect.bisect_left(ones_indices, target)
        # 确认target确实在列表中（避免重复置0操作）
        if pos < len(ones_indices) and ones_indices[pos] == target:
            ones_indices.pop(pos)
    else:
        k = val
        if k > len(ones_indices):
            print(-1)
        else:
            # 第k个1的原索引是ones_indices[k-1]，转换为1-based输出
            print(ones_indices[k-1] + 1)

这个修复后的方案，查询操作是O(logm)（m是当前1的数量），但删除操作是O(m)（因为列表中间删除元素需要移动后续元素），如果你的操作中类型0操作很多，且n很大（比如1e5级别），这个方法还是会超时。

更高效的解决方案：树状数组（Fenwick Tree） #

对于这种单点更新和前缀查询+二分查找的场景，树状数组是最优选择之一，所有操作的时间复杂度都是O(logn)，适合大规模数据。

思路 #

• 树状数组维护每个位置的数值（初始为1，置0时改为0），支持快速计算前缀和（即前i个元素中1的数量）。
• 类型0操作：将对应位置的数值从1改为0，等价于在树状数组中给该位置减去1。
• 类型1操作：通过二分查找找到最小的pos，使得前pos个元素的前缀和等于k（如果总前缀和小于k，输出-1）。

代码实现 #

class FenwickTree:
    def __init__(self, size):
        self.n = size
        self.tree = [0]*(self.n + 1)  # 1-based索引
        # 初始化树状数组，每个位置为1
        for i in range(1, self.n+1):
            self.update(i, 1)
    
    def update(self, idx, delta):
        # 更新idx位置（1-based），加上delta
        while idx <= self.n:
            self.tree[idx] += delta
            idx += idx & -idx
    
    def query(self, idx):
        # 查询前idx个元素的前缀和（1-based）
        res = 0
        while idx > 0:
            res += self.tree[idx]
            idx -= idx & -idx
        return res

t = int(input())
q = int(input())
ft = FenwickTree(t)

for _ in range(q):
    op, val = map(int, input().split())
    if op == 0:
        # 类型0：将第val个位置（1-based）置为0，即减去1
        current = ft.query(val) - ft.query(val-1)
        if current == 1:  # 避免重复置0
            ft.update(val, -1)
    else:
        k = val
        total = ft.query(t)
        if k > total:
            print(-1)
            continue
        # 二分查找第k个1的位置
        left, right = 1, t
        ans = -1
        while left <= right:
            mid = (left + right) // 2
            cnt = ft.query(mid)
            if cnt >= k:
                ans = mid
                right = mid -1
            else:
                left = mid +1
        print(ans)

另一种选择：线段树 #

线段树也可以高效处理这个问题，每个节点维护对应区间内1的数量。类型0操作是单点更新，类型1操作是直接查询第k个1的位置（不需要额外二分，线段树内部可以通过比较左子树的数量来导航）。

代码实现（简化版） #

class SegmentTree:
    def __init__(self, size):
        self.n = size
        self.size = 1
        while self.size < self.n:
            self.size <<=1
        self.tree = [0]*(2*self.size)
        # 初始化叶子节点
        for i in range(self.n):
            self.tree[self.size + i] =1
        # 构建树
        for i in range(self.size-1, 0, -1):
            self.tree[i] = self.tree[2*i] + self.tree[2*i+1]
    
    def update(self, pos, val):
        # pos是0-based原索引
        pos += self.size
        self.tree[pos] = val
        pos >>=1
        while pos >=1:
            new_val = self.tree[2*pos] + self.tree[2*pos+1]
            if self.tree[pos] == new_val:
                break
            self.tree[pos] = new_val
            pos >>=1
    
    def find_kth(self, k):
        # 找第k个1的位置（0-based），不存在返回-1
        if self.tree[1] <k:
            return -1
        node =1
        while node < self.size:
            left = self.tree[2*node]
            if left >=k:
                node =2*node
            else:
                k -= left
                node =2*node+1
        return node - self.size

t = int(input())
q = int(input())
st = SegmentTree(t)

for _ in range(q):
    op, val = map(int, input().split())
    if op ==0:
        # val是1-based索引，转换为0-based
        idx = val-1
        if st.tree[st.size + idx] ==1:
            st.update(idx, 0)
    else:
        k = val
        res = st.find_kth(k)
        print(res+1 if res !=-1 else -1)

这两种数据结构都能在O(logn)的时间内完成每个操作，适合处理大规模的输入数据。

内容的提问来源于stack exchange，提问作者Om Sharma