邁向王者的旅途: [筆記] Codeforces 1208D. Restore Permutation

[題意] 給一串 n 個數字的序列

$s_1, s_2, s_3, ... , s_n$ ，要找出 1 ~ n 這 n 個正整數的排列

$p_1, p_2, ... , p_n$ 使得

$s_i = \sum_{j=1}^{i} p_j$ ，其中

$p_j < p_i$
這題的思考方向是要由後往前看，原因在於

由後往前看時，在你後面已經決定好的數值不會影響你現在正要決定的數字。舉例來說假設 n = 5，並且你已經算出 $p_5$ ，正要決定 $p_4$ 是哪個數字，這時 $p_5$ 並不會影響你 $p_4$ 應該要是多少 (除非你 $p_5$ 是錯的)
由後往前看時，假設要決定 $p_i$ 有好幾個數字都能符合 $s_i$ ，那 $p_i$ 一定是其中。比方說 $p_i$ 有 ${p^1_i, p^2_i, ... , p^k_i}$ ，並且裡面數字最小的是 $p^1_i$ ，如果今天選擇 $p_i = p^2_i$ ，那表示 $s_i$ 必定包含 $p^1_i$ ，那麼 $p^1_i$ 就不該是 $p_i$ 的候選之一，產生矛盾。

第二點就是應該要由後往前看的最重要原因，因為由前往後看的時候不能說選擇最大的一定就對，因為選次大的也不會有問題。

下一個問題是要怎麼根據

$s_i$ 很快地算出

$p_i$ 是多少？從題目的 n 的大小來看，這一步的時間複雜度不能超過

$log\ n$ ，所以如果我們是掃過所有剩餘數字必定會超過時限。因此接下來這步就是解這題最核心也最不直覺的一步：

首先對任意一個數字 x，從前面的分析我們可以知道所有比 x 小的數字除非出現在 x 後面，不然必定包含在

$s_i$ 中，而出現在 x 後面的表示我們必然已經處理過了，所以一開始就可以排除了，所以如果在下圖例子中來看 (灰底方格是已經被選走的數字，? 是還未被決定的數字)，我們可以知道從左邊數來第 3 個 ? 只有 4 吻合。

所以這題的一個重點在於我們必須要能快速地算出對於一個給定的 x，目前 1 ~ n 中還沒被選走的數字裡又比 x 小的數字總和是多少？而要做到這件事情有 2 個資料結構可以用：

Binary Indexed Tree (BIT) 或者是 Segment Tree (ST)

這兩個資料結構對於我們的需求

都非常吻合。資料結構細節及運作方式不多提，google 很好找的範例。

至於要怎麼找到 x，流程如下

找一最大的 $k$ 使得 $2^k \ \leq \ n$
令 $ans = x + 2^k$ ( $x$ 初始值為 0)
用 BIT 或 ST 找 $ans$ 的 prefix sum (就是比 $ans$ 解尚未被拿走的數字總和)，假設是 $BIT[ans]$ (用 ST 就是 $ST[1 ~ ans]$ ，後面都用 BIT 且不再贅述)
若 $BIT[ans] > s_i$ 則 $k = k - 1$ ，回到步驟 2
若 $BIT[ans] < s_i$ ，則 $x = 2^k, k = k - 1$ ，回到步驟 2
若 $BIT[ans] = s_i$ ，則 $ans$ 就是答案並更新 BIT

實作如下 (有點 tricky)：

	/*
	* =====================================================================================
	*
	* Filename: 1208D.cpp
	*
	* Description: codeforces 1208D. Restore Permutation
	*
	* Version: 1.0
	* Created: 2019年09月05日
	* Revision: none
	* Compiler: gcc
	*
	* Author: lionking
	* Organization: None
	*
	* =====================================================================================
	*/

	#include <cstdio>

	using namespace std;

	typedef long long int lli_t;


	// update the value at index 'index' with change 'val' in binary indexed tree
	void update(lli_t *bit, int N, int index, lli_t val)
	{
	for (; index <= N; index += (index & (-index))) {
	bit[index] += val;
	}
	}

	// Given val (s[i] in problem description), find the desired p[i]
	int query(lli_t *bit, int N, int logN, lli_t val)
	{
	int ans = 0;
	for (int i = logN; i >= 0; --i) {
	auto idx = ans + (1 << i);
	if ((idx <= N) && (bit[idx] <= val)) {
	ans = idx;
	val -= bit[idx];
	}
	}

	return (ans + 1);
	}

	int main()
	{
	constexpr int MAXN = 200007;
	// bit: binary indexed tree (BIT)
	lli_t bit[MAXN] = {0}, p[MAXN];
	int N;

	scanf("%d", &N);
	// read s[i] and construct BIT
	for (int i = 1; i <= N; ++i) {
	scanf("%lld", p + i);
	update(bit, N, i, i);
	}

	int logN = 0;
	for (int x = N; x > 0; x >>= 1, ++logN) {}

	// find the answer with the help of BIT from back to front
	for (int i = N; i > 0; --i) {
	const auto ans = query(bit, N, logN, p[i]);
	p[i] = ans;
	update(bit, N, ans, -ans);
	}

	for (int i = 1; i <= N; ++i) {
	printf("%lld ", p[i]);
	}
	putchar('\n');

	return 0;
	}

邁向王者的旅途