#include <bits/stdc++.h>
using namespace std;
vector<long long> P;
vector<int> kpi;
pair<vector<int>, long long> transaction(long long M) {
vector<int> out;
for (int i = 0; i < P.size(); i++) if (M >= P[i]) {
M -= P[i];
kpi[i]--;
out.push_back(i);
}
return {out, M};
}
// start of solution
/*
#include <bits/stdc++.h>
#include "souvenirs.h"
using namespace std;
*/
struct Equation {
vector<int> elem; long long sum;
Equation(vector<int> e, long long s) : elem(e), sum(s) {}
void clear(vector<long long>& P, int bottom) {
while (elem.back() > bottom) { // elem[0] should not be less than bottom, so no need to check elem.empty();
sum -= P[elem.back()];
elem.pop_back();
}
}
};
Equation modified_transaction(long long M) {
pair<vector<int>, long long> raw = transaction(M);
return {raw.first, M - raw.second};
}
void buy_souvenirs(int N, long long P0) {
vector<long long> P(N, 0);
vector<int> kpi(N); // track how many of each type has been bought
for (int i = 0; i < N; i++) kpi[i] = i;
stack<Equation> S; // store info
int bottom = N - 1; // found P[i] forall i > bottom;
S.push({{0}, P0});
// find P
while (bottom) {
Equation u = S.top();
u.clear(P, bottom);
if (u.elem.size() == 1 && u.elem[0] == bottom) {
S.pop();
P[bottom] = u.sum;
bottom--;
} else {
long long amount = (u.sum - 1)/u.elem.size();
Equation next = modified_transaction(amount);
for (auto i: next.elem) kpi[i]--;
S.push(next);
}
}
// clean up
for (int i = 0; i < N; i++) while (kpi[i]--) transaction(P[i]);
}
// end of solution
bool test(vector<long long> T) {
P = T;
int N = P.size();
kpi.resize(N); for (int i = 0; i < N; i++) kpi[i] = i;
buy_souvenirs(N, P[0]);
for (auto i: kpi) if (i) return false;
return true;
}
signed main() {
// testing
vector<vector<long long>> tests;
tests.push_back({4, 3, 1});
tests.push_back({4, 3, 2});
tests.push_back({3, 2, 1});
tests.push_back({10, 9});
tests.push_back({11, 10, 8, 6, 4, 3, 2});
for (int i = 0; i < tests.size(); i++) {
cout << "Test #" << i << ": ";
if (test(tests[i])) cout << "Passed";
else cout << "Failed";
cout << "\n";
}
}
