Using method of Lagrange multipliers in an AtCoder problem

Hi all, I recently attempted Square Price. In this problem we are given N items, and it costs $$$k^2P_i$$$ yen to buy k units of the i-th product. We need to maximize $$$\sum_{i = 1}^{n}k_i$$$ given the total cost cannot exceed M.

While the standard approach to solve this is with binary search, my approach was to treat this as a constrained optimization problem. We need to maximize $$$f = \sum_{i = 1}^{n}k_i$$$ under the constraint $$$g = \sum_{i = 1}^{n}k_i^2P_i = M$$$. So for now, let us allow $$$k_i$$$ to range over all positive real numbers.

Setting $$$\nabla f = \lambda\nabla g$$$, we have a system of equations:
$$$\frac{\partial \sum_{j = 1}^{n}k_j}{\partial k_i} = \lambda \frac{\partial \sum_{j = 1}^{n}k_j^2P_j}{\partial k_i}$$$ for all $$$1 \leq i \leq n$$$
$$$\implies 1 = 2\lambda k_{i}P_i$$$
$$$\implies k_i = \frac{1}{2\lambda P_i}$$$

Along with:
$$$\sum_{i = 1}^{n}k_i^2P_i = M$$$
Plugging in the expression for $$$k_i$$$ gives:
$$$\frac{1}{4\lambda^2}\sum_{i = 1}^{n}\frac{1}{P_i} = M$$$
$$$\implies \lambda = \frac{1}{2}\sqrt{\frac{\sum_{i = 1}^{n}\frac{1}{P_i}}{M}}$$$

Finally we get the expression for $$$k_i$$$
$$$k_i = \frac{1}{P_i \sqrt{\frac{\sum_{j = 1}^{n}\frac{1}{P_j}}{M}}}$$$

Now consider our original problem which requires $$${k_i}$$$ to be integers. The answer must be one of the integer coordinate points in N-dimensional obtained by either rounding down or rounding up each $$$k_i$$$. However we do not need to check all of them, we can simply round each $$$k_i$$$ down, and then greedily round up the coordinate which contributes the least cost, i.e. with the least value of $$$(2k_i + 1)P_i$$$, while the total cost does not exceed M. This works because rounding up any coordinate contributes exactly 1 to f, thus it is never better to pick a more expensive rounding up.

Here is my code:

Code

#include <bits/stdc++.h>
using namespace std;

#define int long long
#define endl "\n"

signed main()
{
    cin.tie(0)->sync_with_stdio(false);
    int n, m;
    cin >> n >> m;
    vector<int> p(n);
    for(auto &i : p)
        cin >> i;
    __float128 rootm = sqrt(m);
    __float128 rootsum = 0;
    for(int i = 0; i < n; i++)
        rootsum += ((__float128)1/((__float128)p[i]));
    rootsum = sqrt((long double)rootsum);
    vector<int> k(n);
    for(int i = 0; i < n; i++)
        k[i] = rootm/(rootsum*p[i]);
    int M = m;
    for(int i = 0; i < n; i++)
        M -= k[i]*k[i]*p[i];
    vector<int> inds(n);
    iota(inds.begin(), inds.end(), 0);
    sort(inds.begin(), inds.end(), [&](int x, int y){return (2*k[x] + 1)*p[x] < (2*k[y] + 1)*p[y];});
    int ans = 0;
    for(auto i : inds)
    {
        if((2*k[i] + 1)*p[i] <= M)
        {
            ans++;
            M -= (2*k[i] + 1)*p[i];
        }
        else
            break;
    }
    for(int i = 0; i < n; i++)
        ans += k[i];
    cout << ans << endl;
}

Unfortunately, I do not get AC with this approach.

Spoiler