codeforces-1208D Restore Permutation题解

发表于 更新于 分类于 阅读次数：

描述

对于[1, n]闭区间的所有整数的某个permutation permutation, 我们有数组nums. 定义nums的每个元素nums[i]表示在permutation的下标i之前的比permutation[i]小的所有数的和.
例如如果nums是0 0 0, 那么permutation是[3,2,1]. 因为对于每个数, 没有在它之前的比它小的数. 如果nums是0 1 1 1 10, permutation是[1,4,3,2,5].

思路

考虑nums数组的最后一个元素. 因为这是数组的最后, 所有比它(放在这个位置的[1,n]内的某个数字)小的数字都会出现在它前面. (不可能出现在后面因为这已经是最后了). 假设这个最后应该放置的数字是k, 那么这些比它的小的数字之和必然是sum([1...k-1]).
对于上面的式子可能看做是一种前缀和. 所以通过搜索(这个前缀和在所有前缀和中的位置)可以找到到k的值.

比如, 对于[1,4,2,3]这个permutation来说, nums是[0,1,1,3]. 最后一个值是3,由1+2得到, 那么放在这个位置的数一定是1,2后面的第一个数字3. 拿到3之后, 剩余的数字1,2,4起始也可以看做一个permutation, 只要有序无重复就行.(它仍然拥有一个对应的前缀和数组). 但是上一个前缀和数组含有3的影响, 现在已经找到了3,可以把它删掉(置为0).

需要动态维护这样一个前缀和数组, 使用了线段树(权值线段树). 每个节点存的是[l, r]范围内的和. 首先把[1,n]所有数插入进去, 然后query当前nums数组的最后一个值. 返回的值就是应该放在这个位置的数字, 然后把这个数字置为0. 需要注意的是, 这样我们的前缀和数组的生成就会有很多0的影响, eg. [1,3,3,3,3,3,6], 因为原始数组出现了多个0, 前缀和数组会有重复数字. 这个时候再查找目标前缀和, 我们需要找到尽量右侧的点(坐标).类似bisect.bisect_right().

所以实际query时, 参数是nums[i]+1, 保证不受0的干扰. (其实还有一种情况, [1,2,3]的前缀和是[1,3,6]) 搜6时会落到3上, 但和6对应的是4.

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class segTree():
    def __init__(self, n):
        self.t = [0] * (n << 2)
 
    def update(self, node, l, r, index, value):
        if l == r:
            self.t[node] = value
            return
        mid = (l + r) >> 1
        if index <= mid:
            self.update(node*2, l, mid, index, value)
        else:
            self.update(node*2 + 1, mid + 1, r, index, value)
        self.t[node] = self.t[node*2] + self.t[node*2 + 1]    # 更新完两个子节点还要计算一下当前节点
 
    def query(self, node, l, r, value):
        if l == r:
            return self.t[node]
        mid = (l + r) >> 1
        if self.t[node*2] >= value:
            return self.query(node*2, l, mid, value)
        return self.query(node*2 + 1, mid + 1, r, value - self.t[node*2])   # 向右搜索时, 减掉左侧节点的值, 这样最终会落到目标点上
 
def do():
    n = int(input())
    nums = [int(i) for i in input().split(" ")]
    res = [0]*n
    weightTree = segTree(n)
    for i in range(1, n+1):
        weightTree.update(1, 1, n, i, i)
    # print(weightTree.t)
    for i in range(n-1, -1, -1):
        res[i] = weightTree.query(1, 1, n, nums[i] + 1)
        weightTree.update(1, 1, n, res[i], 0)
    return " ".join([str(c) for c in res])
 
print(do())