Here's my attempt at solving the problem 1829G
I've started from the bottom nth row and worked my way up from there. The set considered
has all the values that have already been taken before. This solution fails for the sample test case when n = 1434
( gives 80385350
, should give 63145186
)
Code
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define mod 1000000007
#define nl "\n"
#define F first
#define S second
#define no cout << "NO" << nl
#define yes cout << "YES" << nl
#define all(x) x.begin(), x.end()
#define vll vector<ll>
#define fo(i, n) for (ll i = 0; i < n; ++i)
#define fo1(i, n) for (ll i = 1; i < n; ++i)
#define show(x) cout << #x << " = " << x << nl
ll n;
void solve()
{
cin >> n;
// if cur_row = r;
// then next contribution for row r-1 is cur_el - r, cur_el-r+1 squared respectively
vector<ll> sums;
ll i = 1, j = 1;
ll curSum = 1;
while (i <= 1e7)
{
sums.push_back(curSum);
j++;
i += j;
curSum = (i * (i + 1)) / 2;
}
set<ll> cur, considered;
ll ans = 0;
cur.insert(n);
considered.insert(n);
while (!cur.empty())
{
for (const ll &el : cur)
{
ans += el * el;
}
set<ll> prev;
for (const ll &el : cur)
{
ll idx = lower_bound(all(sums), (el * (el + 1) / 2)) - sums.begin();
if (idx != sums.size())
{
if (sums[idx] != (el * (el + 1) / 2))
{
ll row = idx + 1;
if (el - row > 0 && considered.find(el - row) == considered.end())
{
prev.insert(el - row);
considered.insert(el - row);
}
if (el - row + 1 > 0 && considered.find(el - row + 1) == considered.end())
{
prev.insert(el - row + 1);
considered.insert(el - row + 1);
}
}
else
{
ll row = idx+1;
if (el - row > 0 && considered.find(el - row) == considered.end())
{
prev.insert(el - row);
considered.insert(el - row);
}
}
}
}
cur = prev;
}
cout << ans << "\n";
}
int main()
{
ios_base::sync_with_stdio(false);
cin.tie(NULL), cout.tie(NULL);
ll _;
//_=1;
cin >> _;
while (_--)
{
solve();
}
return 0;
}
Where am I messing up here?