오랜만에 한 문제 풀어봤다.
Problem - E - Codeforces
codeforces.com
어떤 왕국이 n개의 정점을 가지는 트리 구조라고 하자. 각각의 정점은 a_i라는 값을 가지고 있다. 모든 정점은 일자로 연결되어 있다. 정확하게, 1부터 n까지 모든 i에 대해서, i번째 정점과 i+1번째 정점은 간선으로 연결되어 있다.
여기서 함수 f(l, r)을 정의한다. (l <= r)
* 먼저, a_i가 l보다 크거나 같고 r보다 작거나 같은 정점만 남긴다.
* 위 조건을 만족하는 정점으로만 표현된 그래프에서 connected component의 개수를 f(l, r)로 정의한다.
우리가 해야할 것은 이런 함수들의 합을 구하는 것이다.
입력
첫 번째 줄은 정점의 개수 n을 나타낸다. (1<=n<=100000)
두 번째 줄은 각 정점의 a_i 값을 나타낸다. (1<=a_i<=n)
출력
함수의 합을 하나의 숫자로 나타낸다.
처음에는 union and find로 O(nlogn)으로 푸는 건줄 알았는데, 잘 보면 f(l, r)을 O(1)에 구해도 합을 구해야 하기 때문에 O(n^2)가 되어 시간이 초과되어 버린다.
이런 경우, 두 시그마 중에 하나를 O(logn)이나 O(1)으로 내려 전체 시간 복잡도를 O(n)으로 낮춰야 된다.
난 똑똑하지 않아서,, 그냥 해설 봤다. (뭔가 4시간 고민하고 결국 해설 킬 것 같은 강한 느낌이 오더라..)
해설을 한글로 풀고 작성한 소스 코드를 첨부했다.
직접 풀고싶은 사람은 아래 글을 펼치지 말 것!
어쨋든, 전체 시간 복잡도를 O(n) 또는 O(nlogn)으로 줄여야 한다.
아이디어는 하나의 정점이 전체 합에서 기여하는 값 (Contribution of sum)을 구하는 것이다.
그러기 위해서 먼저, 뭘 기여하는 지 생각해봐야 한다.
각각의 함수 값이 connected component의 개수를 의미하므로, 만약 이 정점으로 인해서 connected component가 생성된다면, 그 정점이 1만큼 기여를 했다고 생각할 수 있다.
먼저, 정점들이 일자로 연결되어 있기 때문에 그래프에 포함된 연속적인 정점들이 하나의 connected component를 구성한다고 생각하면 된다.
여러 정점이 모여서 하나의 connected component를 만들기 때문에, connected component에서 가장 작은 index를 가진 정점이 connected component를 만드는 데 기여를 했다고 생각하면 편하다.
그리고 이제 각각의 정점이 모든 함수에 대해서 몇 번 기여를 하는 지를 수학적인 식으로 나타내보자.
i번째 정점의 기여도를 계산하기 위해서 보아야 할 것은, i-1번째 정점의 기여도이다.
만약, "i-1번째 정점의 값이 i번째 정점의 값보다 작거나 같을 경우"를 생각해보자.
l이 i-1번째 정점의 값보다 작다면, 그래프에 i-1번째 정점이 포함되게 된다.
이 경우, r이 i번째 정점의 값보다 작다면 그래프에 포함되지 않아서 기여를 하지 못한다.
또한, r이 i번째 정점의 값보다 크거나 같다면 i번째 정점이 그래프에 포함되지만, i-1번째 정점이 이미 그래프에 포함되어 있기 때문에 connected component에서 가장 작은 index가 되지 못하기 때문에 기여를 하지 못하게 된다.
즉, i번째 정점이 connected component에서 가장 작은 index가 되어서 함수값에 기여를 하기 위해서는 i-1번째 정점이 그래프에 포함되지 않아야 한다는 것을 알 수 있다.
l이 i-1번째 정점의 값보다 크다면, 그래프에 i-1번째 정점이 포함되지 않게 된다.
이 경우, r이 i번째 정점의 값보다 작다면, 마찬가지로 그래프에 포함되지 않아서 기여를 하지 못하게 된다.
하지만, r이 i번째 정점의 값보다 크거나 같다면 i번째 정점이 그래프에 포함되면서, connected component에서 가장 작은 index로 동작하게 되어 기여를 할 수 있게 된다.
여기서 l이 i번째 정점의 값보다 커지면 또 다시 i번째 정점이 그래프에서 빠지게 되므로 l은 i번째 정점보다 작거나 같아야 된다.
즉, i번째 정점이 기여를 할 수 있는 상황은 (1) l이 i-1번째 정점의 값보다 크고, (2) r이 i번째 정점의 값보다 크거나 같아야 하면서 (3) l이 i번째 정점의 값보다 작거나 같을 때 반드시 기여를 할 수 있다.
이를 기반으로 전체 경우에서 기여하는 경우의 수를 구해보면 다음고 ㅏ같다.
먼저 가능한 r은 i번째 정점의 값부터 전체 노드의 개수 n까지 이다.
그리고 각각의 r에 대해서 가능한 l은 i-1번째 정점의 값부터 i번째 정점의 값까지이다.
따라서 i번째 정점의 contribution of sum은 (r-a_i+1) * (a_{i}-a{i-1})이 된다.
여기까지가 "i-1번째 정점의 값이 i번째 정점의 값보다 작거나 같을 경우"일 때의 시나리오이다.
거의 똑같은 로직으로 "i-1번째 정점의 값이 i번째 정점의 값보다 클 경우"일 때의 시나리오를 계산할 수 있다.
계산하면, i번째 정점의 contribution of sum은 (a_i) * (a_{i-1}-a{i})가 된다.
마지막으로 처리해야 할 문제는 1번째 정점을 위한 "가상의 0번째 정점의 값 설정"이다.
결론적으로 답은 a_0 = 0이다.
왜냐하면, a_1은 항상 connected component에서 가장 작은 index로 작용하기 때문에 l이 a_1보다 작거나 같고 r이 a_1보다 크거나 같으면 반드시 기여를 하고, 아닐 경우, 기여를 하지 못한다. 즉, a_1의 contribution of sum은 (r-a_1+1)*(a_1-0)가 되게 된다.
a_0 = 0일 경우, "i-1번째 정점의 값이 i번째 정점의 값보다 작거나 같을 경우"일 때 contribution of sum과 같은 꼴이 되고 a_0를 0으로 두기 때문에 1<=a_i라는 조건에 의해 0번째 정점의 값이 1번째 정점의 값보다 항상 작아서 모든 조건을 자연스럽게 만족하게 된다.
난 이제 이런 문제 생각해서 풀수나 있을까 싶다. 이산수학 문제인데, 해설 자체는 간단한데 이걸 만들어내기까지가 참 어려운 것 같다. 아래는 소스 코드다.
/* Contribution of sum */
#include<iostream>
#include<vector>
using namespace std;
typedef long long int ll;
int main(){
ll n;
ll ans = 0;
cin >> n;
vector<ll> in(n+1);
for(int i=1; i<=n; i++)
cin >> in[i];
in[0] = 0;
for(int i=1; i<=n; i++){
// 만약 a_{i-1} <= a_{i}일 경우,
// l <= a_{i-1}일 때는, a_{i-1}이 connected component를
// 형성하기 때문에 a_{i}의 contribution이 0임.
// l > a_{i-1}일 때는, a_{i-1}이 graph에 포함되지 않기 때문에
// a_{i}가 connected component를 만들게 됨.
// l > a_{i-1}이면서 a_{i} <= r일 때, a_{i}가 graph에
// 포함되기 때문에 위의 조건이 되었을 때, contribution이 1임.
// 가능한 r은 [a_{i}, n]이고,
// 각각의 r에 대해 가능한 l은 [a_{i-1}+1, a_{i}] 이다.
if(in[i-1] <= in[i])
ans += (n-in[i]+1) * (in[i]-in[i-1]);
// 만약 a_{i-1} > a_{i}일 경우,
// r < a_{i-1}일 때는, a_{i}가 connected component를
// 형성하기 때문에 a_{i}의 contribution이 1임.
// r >= a_{i-1}부터, a_{i-1}이 connected component를
// 형성하기 때문에 a_{i}의 contribution은 0이 됨.
// l <= a_{i}이면서 r < a_{i-1}일 때, a_{i}가 graph에
// 포함되기 때문에 위의 조건이 되었을 때, contribution이 1임.
// 가능한 l은 [1, a_{i}]이고,
// 각각의 l에 대해 가능한 r은 [a_{i}, a_{i-1}-1]이다.
else
ans += in[i] * (in[i-1]-in[i]);
}
cout << ans;
return 0;
}