[SWEA] 재성이의 비료주기

[SWEA] 재성이의 비료주기

Problem URL : 재성이의 비료주기

#include <vector>
#include <cstring>

using namespace std;

typedef long long ll;

ll sum(vector<ll>& tree, int i){
	ll ans = 0;
	while (i > 0){
		ans += tree[i];
		i -= (i & -i); 
	}

	return ans;
}

void update(vector<ll>& tree, int i, ll diff){	
	while (i < tree.size()){		
		tree[i] += diff;
		i += (i & -i);
	}
}

int main() {
	int TC;
	scanf("%d", &TC);
	for (int tc = 1; tc <= TC; tc++) {
		int N, D;
		long long s, a, b;
		scanf("%d %d %lld %lld %lld", &N, &D, &s, &a, &b);
		vector<ll> arr(N + 1);
		vector<ll> tree(N + 1);
		fill(arr.begin(), arr.end(), 0);
		fill(tree.begin(), tree.end(), 0);
		ll ans = 0;
		int idx = 1;
		for (int i = 1; i <= D; i++) {
			int c = s % N + 1;
			s = (s * a + b) % 1000000007;
			int k = s % N + 1;
			s = (s * a + b) % 1000000007;
			update(tree, idx, i);
			idx += c;
			if (idx > N) {
				idx -= N;
				if (idx > 1) {
					update(tree, 1, i);
					update(tree, idx, -i);
				}
			} else {
				update(tree, idx, i * (-1));
			}
			
			int find = idx + k - 1;
			if (find > N) {
				find -= N;
			}
			ans += sum(tree, find);
		}
		printf("#%d %lld\n",tc,ans);
	}
}

Comments

링크 를 참고했다.
펜윅트리는 펜윅트리 1, 펜윅트리 2 를 참고하자

틀린 풀이(제한 시간 초과)

#include <vector>
#include <cstring>
#include <iostream>

using namespace std;

typedef long long ll;

int arr[100001];

int main() {
	int TC;
	scanf("%d", &TC);
	for (int tc = 1; tc <= TC; tc++) {
		int N, D;
		long long s, a, b;
		scanf("%d %d %lld %lld %lld", &N, &D, &s, &a, &b);
		
		ll ans = 0;
		int idx = 1;
		memset(arr, 0, sizeof(arr));
		for (int i = 1; i <= D; i++) {
			int c = s % N + 1;
			s = (s * a + b) % 1000000007;
			int k = s % N + 1;
			s = (s * a + b) % 1000000007;
			
			if (idx + c - 1 > N) {
				for (int j = idx; j <= N; j++) {
					arr[j] += i;
				}
				for (int j = 1; j <= idx + c - 1 - N; j++) {
					arr[j] += i;
				}
				idx += c - N;
			} else {
				for (int j = idx; j <= idx + c - 1; j++) {
					arr[j] += i;
				}
				idx += c;
			}

			k += idx - 1;
			if (k > N) {
				k -= N;
			}
			
			ans += arr[k];
		}
		printf("#%d %lld\n",tc,ans);
	}
}