片側にある自分より大きい数をO(N)で求める方法
こんにちは。tatsyです。
なんとも分かりづらいタイトルですが、この記事ではある配列が与えられたときに、各要素に対して自分より右(もしくは左)にある自分より大きいもののうち一番近いものを割り当てる問題を解きます。
やっぱり分かりづらい・・・つまり、例を示すなら、右にある自分より小さい数を見つけるとして、
入力: [4, 2, 3, 5, 1]
出力: [2, 1, 1, 1, -1]
みたいになります。一番右の1は自分より小さい数がもう右にないので-1を入れています。
問題自体はとても簡単で各要素に対して、自分より右にあるものを1つずつ調べていけばよいのですが、これだと計算量はO(N2)です。
アルゴリズム
実はこの問題にはO(N)で解くアルゴリズムが存在していて、C++で書くと次のような関数になります。
v) {
int n = (int)v.size();
vector<int> res(n, -1);
stack<int> s;
for(int i=0; i<n; i++) {
while(!s.empty() && v[i] < v[s.top()]) {
res[s.top()] = v[i];
s.pop();
}
s.push(i);
}
return res;
}
このアルゴリズムはstackに「まだ自分より小さい数が見つかっていない要素」を積んでおいて、for文を回して、小さい値が現れたら結果の配列を更新するという方法です。
適当に配列を与えて試してもらえれば、正しそうなことが分かるはずです。が、果たしてこれはO(N)なのか？という疑問は残ります。
for文がN回繰り返しなので、ここは良いのですが、for文の中にwhile文があるので、ここの計算量を考えなくてはいけません。
ただ、よく見ると、while文は最大でもstackに追加された要素の数しか繰り返しをしないことが分かります。
さらによく見ると、stackに追加される要素は配列のインデックスでどんなに多くても配列の長さ以上にはなりません。
ということはwhile文の中身が実行される回数はプログラム全体で高々N回なので、結局、計算量はO(N)です。
練習用の問題
このアルゴリズムを使う問題がCodeforcesにあります。
僕もこの問題とEditorialを読んでいて今回のアルゴリズムを見つけました。
問題はよく分からないシチュエーションですが、とりあえずN人の超能力者がいます。
各超能力者は強さみたいなものを持っており、自分より右側にいて、かつ自分より弱い超能力者を消すことができます。
1ターンに最大1人しか消せないと考えたとき、それ以上状態が変化しなくなるのは何ターン後かを求めます。
やり方を説明する前に、各超能力者が消されるまでの時間をtiとしましょう。
すると、ある超能力者iが消されるまでの時間は、その超能力者より左にいて、かつ自分より強い超能力者の番号をj (j<i) としたときに、
$$ t_i = max(t_{j+1}, ldots , t_{i-1}) $$
となるはずです。
ですから、自分より左にいて、自分より強い超能力者を上のアルゴリズムで見つけながら、各超能力者の消されるまでの時間を求めます。
消されてしまう超能力者の生存時間の中で最も大きなものが答えです。
以下が正解のコードになります。
#include <sstream>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <ctime>
#include <cstring>
#include <climits>
#include <algorithm>
#include <map>
#include <set>
#include <stack>
#include <vector>
#include <queue>
#include <bitset>
#include <string>
using namespace std;
#define REP(i,n) for(int i=0; i<n; i++)
#define RPA(i,s,e) for(int i=s; i<=e; i++)
#define RPD(i,s,e) for(int i=s; i>=e; i--)
#define PB(a) push_back(a)
#define MP(i,s) make_pair(i,s)
#define SZ(a) (int)(a).size()
#define ALL(a) (a).begin(), (a).end()
#define PRT(a) cout << (a) << endl
#define PR2(a,b) cout << (a) << " " << (b) << endl
#define PR3(a,b,c) cout << (a) << " " << (b) << " " << (c) << endl
typedef vector<int> VI;
typedef long long LL;
typedef pair<int,int> P;
const int INF = 1 << 30;
int n;
int idx[100010];
int life[100010];
int main() {
while(scanf("%d", &n)!=EOF) {
idx[0] = 0;
RPA(i,1,n) {
scanf("%d", &idx[i]);
}
life[0] = INF;
stack<int> s;
s.push(0);
RPA(i,1,n) {
life[i] = 1;
while(!s.empty() && idx[i] >= idx[s.top()]) {
life[i] = max(life[i], life[s.top()]+1);
s.pop();
}
s.push(i);
}
int ans = 0;
RPA(i,1,n) {
if(life[i] < INF) {
ans = max(ans, life[i]);
}
}
cout << ans << endl;
}
}
この問題は、入力の配列サイズが105なので、O(N2)のアルゴリズムだとTLEになってしまいます。
ていうか、このO(N)アルゴリズムって常識なのだろうか？
またひとつ賢くなりました、というところで今回の記事を終わります。
最後までお読みいただきありがとうございました！