AcWing 244. 谜一样的牛
原题链接
简单
作者:
皓首不倦
,
2020-08-27 21:33:37
,
所有人可见
,
阅读 314
'''
最后一一头牛的高度明显是一开始就可以推断出来的,假设其高度是h(n-1)
接着从1-n这n个数字中把h(n-1)删掉,然后再看倒数第二头牛,假设其前面
有vals[n-2]头牛比它矮,说明这头牛是剩下的数里面第vals[n-2]+1大的数
假设exist[i]为1表示高度i还没有被选走,0表示已经被选走,那要找剩下的
高度中,第vals[n-2]+1大的数也就是找最小的x 让exist序列以x位置结尾的
前缀和刚好等于vals[n-2]+1,显然可以对x进行二分搜索,而exist序列由于
在选择了牛的高度后动态变化,又要维护其前缀和,刚好符合树状数组的特性,
用树状数组维护exist序列的前缀和即可,每确定一头牛的高度,就从exist序列
中把对对应位置设置为0即可,这样从后往前进行迭代,每头牛的高度都可以二分
搜索搜出来
'''
import sys
class FenwickTree:
def __init__(self, data):
if len(data) == 0:
raise ValueError('data length is zero!')
self.n = len(data)
prefix_sum = [0] * (self.n + 1)
self.sum_seg = [0] * (self.n + 1) # 分段存储的区间和 sum_seg[x]表示的是x位置结尾,长度为x&(-x)的区间中的数值的和
# 实际存储时候错一位存储,外部数据的下标从0开始,但是树状数组的下标是从1开始的
i = 1
for val in data:
prefix_sum[i] = prefix_sum[i-1] + val
self.sum_seg[i] = prefix_sum[i] - prefix_sum[i & (i-1)]
i += 1
# 获取x位置的原始数值
def getOrigValue(self, x):
x += 1
if x < 1 or x > self.n:
raise IndexError(f'error idx = {x}')
ans = self.sum_seg[x]
lca = x & (x-1) # 最近公共祖先
x -= 1
while x != lca:
ans -= self.sum_seg[x]
x &= (x-1)
return ans
# 获取区间[0, end]的数值和
def __getSumWithEnd(self, end): # 这里end是内部坐标,不是外部坐标
ans = 0
while end:
ans += self.sum_seg[end]
end &= (end - 1)
return ans
# 获取区间的数据和
def getSum(self, start, end):
start, end = start+1, end+1
if not (end >= start and start >= 1 and end <= self.n):
raise IndexError(f'bad range {(start-1, end-1)}')
if start == 1:
return self.__getSumWithEnd(end)
return self.__getSumWithEnd(end) - self.__getSumWithEnd(start-1)
# 更新x位置数值
def updateValue(self, x, val):
orig_val = self.getOrigValue(x)
x += 1
delta = val - orig_val
while x <= self.n:
self.sum_seg[x] += delta
x += x & (-x)
# x位置增加数值delta
def addValue(self, x, delta):
x += 1
while x <= self.n:
self.sum_seg[x] += delta
x += x & (-x)
n = int(input())
vals = [0]*n # val[x]表示x位置的数值的前面前面比其小的数字的个数
exist = [1]*(n+1) # exist[val] 表示数值val是否还存在没有被选走
exist[0] = 0
for i in range(1, n):
vals[i] = int(input())
tree = FenwickTree(exist)
ans_arr = [0] * n
# 从后往前递推每一个位置的数值
for i in range(n-1, -1, -1):
# 二分搜索,在剩下的数值中找第vals[i]+1大的数值,也就是找最小的x 让tree.sum(0, x) == vals[i]+1
l, r = 1, n
ans = -1
while l <= r:
mid = l + (r-l) // 2
if tree.getSum(0, mid) == vals[i] + 1:
ans = mid
r = mid - 1
elif tree.getSum(0, mid) > vals[i] + 1:
r = mid-1
else:
l = mid + 1
ans_arr[i] = ans
tree.addValue(ans, -1) # 把牛删掉
for val in ans_arr:
print(val)