malteklaes/cpp-code-code_search_net-style huggingface dataset:
#include <bits/stdc++.h>   using namespace std;   long long int digSum(long long int n) {     long long int sum = 0;     while (n != 0) {       sum += n % 10;       n = n / 10;     }     return sum;   }   bool cmp1(int a, int b) { return a < b; }   bool cmp2(int a, int b) { return a > b; }   int main() {     ios_base::sync_with_stdio(false);     cin.tie(NULL);     long long int n, k, sum2 = 0, sum1 = 0, x;     cin >> n >> k;     vector<int> v(n), ans;     for (int(i) = 0; (i) < (n); (i)++) {       cin >> v[i];       sum1 += v[i];     }     x = sum1 / k;     bool ok = true;     if (sum1 % k == 0) {       int i = 0;       while (i < n) {         long long int sum2 = 0;         long long int ct = 0;         while (i < n && sum2 < x) {           sum2 += v[i];           i++;           ct++;         }         if (sum2 == x) {           ans.push_back(ct);         } else {           ok = false;           break;         }       }       if (ok) {         cout <<  Yes ;         cout << endl;         for (int i = 0; i < ans.size(); i++) {           cout << ans[i] <<    ;         }       } else {         cout <<  No ;         cout << endl;       }     } else {       cout <<  No ;       cout << endl;     }     return 0;   }
#include <bits/stdc++.h>   using namespace std;   int p, n;   vector<int> gra[100010];   double dp[100010], sz[100010];   void dfs(int u) {     sz[u] = 1;     int m = gra[u].size();     for (int i = 0; i <= m - 1; i++) {       dfs(gra[u][i]);       sz[u] += sz[gra[u][i]];     }   }   void solve(int u) {     int m = gra[u].size();     for (int i = 0; i <= m - 1; i++) {       dp[gra[u][i]] = dp[u] + 1 + (sz[u] - sz[gra[u][i]] - 1) / 2.0;       solve(gra[u][i]);     }   }   int main() {     dp[1] = 1;     scanf( %d , &n);     for (int i = 2; i <= n; i++) {       scanf( %d , &p);       gra[p].push_back(i);     }     dfs(1);     solve(1);     for (int i = 1; i <= n; i++) printf( %.6f  , dp[i]);     return 0;   }
#include <bits/stdc++.h>   using namespace std;   int main() {     int n, T, low, mid, high, a[200005], t[200005];     vector<int> ans;     scanf( %d%d , &n, &T);     for (int i = 1; i <= n; i++) {       scanf( %d%d , &a[i], &t[i]);     }     low = 0;     high = n;     while (low <= high) {       mid = (low + high) / 2;       vector<pair<int, int> > v;       int cnt = 0, tot_t = 0, idx = -1;       for (int i = 1; i <= n; i++) {         if (a[i] >= mid) v.push_back(make_pair(t[i], i));       }       sort(v.begin(), v.end());       for (int i = 0; i < v.size(); i++) {         if (tot_t + v[i].first > T) break;         cnt++;         tot_t += v[i].first;         if (cnt == mid) {           idx = i;           break;         }       }       if (cnt >= mid) {         ans.clear();         for (int i = 0; i <= idx; i++) ans.push_back(v[i].second);         low = mid + 1;       } else {         high = mid - 1;       }     }     printf( %d n%d n , ans.size(), ans.size());     for (int i = 0; i < ans.size(); i++) printf( %d  , ans[i]);     printf(  n );     return 0;   }
#include <bits/stdc++.h>   #pragma GCC optimize( Ofast )   #pragma GCC target(         avx,avx2,fma,sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,tune=native )   using namespace std;   long long MOD = 998244353;   long long INF = 1e18;   const int N = 100001;   long long dp[N][101];   long long sm[N];   long long cnt[N][101];   int main() {     ios::sync_with_stdio(0), cin.tie(0), cout.tie(0),         cout << fixed << setprecision(20);     long long n, k, len;     cin >> n >> k >> len;     if (len == 1) {       cout << 0;       return 0;     }     vector<long long> a(n + 1);     sm[0] = 1;     for (int i = 1; i <= n; i++) {       cin >> a[i];       for (int j = 1; j <= k; j++) {         cnt[i][j] = cnt[i - 1][j] + (a[i] == -1 || a[i] == j);         if (a[i] != -1 && j != a[i]) {           cnt[i][j] = 0;         }       }       if (a[i] != -1) {         dp[i][a[i]] += sm[i - 1];         if (cnt[i][a[i]] >= len) dp[i][a[i]] -= sm[i - len] - dp[i - len][a[i]];       } else {         for (int j = 1; j <= k; j++) {           dp[i][j] += sm[i - 1];           if (cnt[i][j] >= len) dp[i][j] -= sm[i - len] - dp[i - len][j];         }       }       for (int j = 1; j <= k; j++) {         dp[i][j] %= MOD;         dp[i][j] += MOD;         dp[i][j] %= MOD;         sm[i] += dp[i][j];         sm[i] %= MOD;       }     }     cout << sm[n];   }
#include <bits/stdc++.h>   using namespace std;   const int MAX_N = 100005;   int n;   int m;   int cnt[MAX_N];   int a[MAX_N];   int main() {     scanf( %d %d , &m, &n);     for (int i = 1; i <= n; i++) {       scanf( %d , &a[i]);     }     int diff = 0;     for (int i = 1; i <= n; i++) {       cnt[a[i]]++;       if (cnt[a[i]] == 1) {         diff++;       }       if (diff == m) {         cout << 1;         for (int j = 1; j <= m; j++) {           cnt[j]--;           if (cnt[j] == 0) {             diff--;           }         }       } else {         cout << 0;       }     }     return 0;   }
#include <bits/stdc++.h>   using namespace std;   int main() {     ios::sync_with_stdio(0);     cin.tie(0);     cout.tie(0);     cerr <<  Time :   << 1000 * ((double)clock()) / (double)CLOCKS_PER_SEC          <<  ms n ;     long long t;     cin >> t;     while (t--) {       long long n, g = 0, s = 0, b = 0, i, gidx, bidx, cidx;       cin >> n;       int flag = 0;       long long p[n];       for (i = 0; i < n; i++) {         cin >> p[i];       }       long long m = n / 2;       cidx = m;       cidx--;       while (p[cidx + 1] == p[cidx]) {         cidx--;         if (cidx < 4) {           cout <<  0 0 0                 <<   n ;           flag = 1;           break;         }       }       if (flag) continue;       gidx = 1;       while (p[gidx] == p[gidx - 1]) {         gidx++;         if (gidx >= cidx) {           cout <<  0 0 0                 <<   n ;           flag = 1;           break;         }       }       if (flag) continue;       bidx = 2 * gidx + 1;       while (p[bidx] == p[bidx - 1]) {         bidx++;         if (bidx >= cidx) {           cout <<  0 0 0                 <<   n ;           flag = 1;           break;         }       }       if (flag) continue;       g = gidx;       b = bidx - gidx;       s = cidx - bidx + 1;       if (g == 0 || b == 0 || s <= 0 || g >= s)         cout <<  0 0 0               <<   n ;       else         cout << g <<     << b <<     << s <<   n ;     }     return 0;   }
#include <bits/stdc++.h>   using namespace std;   vector<int> graph[2001];   vector<int> level[2001];   int visited[2001];   int a, b;   void dfs(int u, int depth) {     visited[u] = 1;     level[depth].push_back(u);     for (int v : graph[u]) {       if (!visited[v]) {         dfs(v, depth + 1);       }     }   }   int query = 0;   void ask(vector<int>& v) {     cout <<  ?   << v.size() <<    ;     for (int x : v) {       cout << x <<    ;     }     cout <<   n  << flush;     cin >> a >> b;   }   void bs(int l, int r, int ref) {     int curr_node = 0, curr_level = 0;     while (l <= r) {       if (r - l <= 1) {         ask(level[r]);         if (b == ref) {           curr_level = r;           curr_node = a;         } else {           ask(level[l]);           curr_node = a;           curr_level = l;         }         break;       }       int mid = (l + r) / 2;       ask(level[mid]);       if (b == ref) {         l = mid;       } else {         r = mid - 1;       }     }     for (int i = 0; i <= 2000; i++) {       visited[i] = 0;       level[i].clear();     }     dfs(curr_node, 1);     ask(level[ref + 1]);     cout <<  !   << curr_node <<     << a <<   n  << flush;     string s;     cin >> s;   }   int main() {     int t;     cin >> t;     while (t--) {       query = 0;       int n, x, y;       cin >> n;       for (int i = 0; i <= 2000; i++) {         graph[i].clear();         level[i].clear();         visited[i] = 0;       }       for (int i = 0; i < n - 1; i++) {         cin >> x >> y;         graph[x].push_back(y);         graph[y].push_back(x);       }       cout <<  ?   << n <<    ;       for (int i = 1; i <= n; i++) {         cout << i <<    ;       }       cout <<   n  << flush;       cin >> x >> y;       dfs(x, 1);       int len = 0;       for (int i = n; i >= 1; i--) {         if (level[i].size() != 0) {           len = i;           break;         }       }       bs(1, len, y);     }   }
#include <bits/stdc++.h>   using namespace std;   const int maxiz = 1000 * 1000 + 100;   int d[maxiz];   int main() {     int n;     long long A, M = 0;     cin >> n >> A;     for (int i = 0; i < n; i++) {       cin >> d[i];       M = (M + d[i] * 1ll) * 1ll;     }     for (int i = 0; i < n; i++) {       long long maxi, mini;       if (A - M + d[i] > 1)         mini = A - (M - d[i]);       else         mini = 1;       if (d[i] < A - n + 1)         maxi = d[i];       else         maxi = A - n + 1;       cout << d[i] - (maxi - mini + 1) <<    ;     }   }
#include <bits/stdc++.h>   using namespace std;   int f[20001][101];   struct node {     int ai, bi;   } a[1001];   bool cmp(node x, node y) { return x.bi > y.bi; }   int main() {     int n, sum1 = 0, sum2 = 0, d = 0;     scanf( %d , &n);     for (int i = 1; i <= n; i++) {       scanf( %d , &a[i].ai);       sum1 += a[i].ai;     }     for (int i = 1; i <= n; i++) {       scanf( %d , &a[i].bi);     }     sort(a + 1, a + n + 1, cmp);     while (sum1 > sum2) sum2 += a[++d].bi;     printf( %d  , d);     memset(f, 128, sizeof(f));     f[0][0] = 0;     for (int i = 1; i <= n; i++) {       for (int j = sum2; j >= a[i].bi; j--) {         for (int k = 1; k <= d; k++) {           f[j][k] = max(f[j][k], f[j - a[i].bi][k - 1] + a[i].ai);         }       }     }     int ans = 0;     for (int i = sum1; i <= sum2; i++) {       ans = max(ans, f[i][d]);     }     printf( %d n , sum1 - ans);     return 0;   }
#include <bits/stdc++.h>   using namespace std;   const int N = 110;   const int M = 10010;   int n, m;   char mat[N][M];   int dir[2] = {-1, 1};   long long solve() {     long long ans = 0;     int x = 0, y = 0, d = 1;     bool isLeft = true, isRight = true;     int pos[2] = {0, 0};     while (true) {       if (!isLeft && !isRight) return -1;       if (x == n - 1) return ans;       if (mat[x + 1][y] ==  . ) {         ans++;         x++;         isLeft = isRight = true;         pos[0] = pos[1] = y;       } else if (y == pos[d]) {         int ty = y + dir[d];         ans++;         if (ty < 0 || ty >= m || mat[x][ty] ==  # ) {           if (d == 0)             isLeft = false;           else             isRight = false;           d ^= 1;         } else {           pos[d] += dir[d];           if (mat[x][ty] ==  . )             y += dir[d];           else             d ^= 1;         }       } else {         ans += (long long)abs((long long)pos[d] - (long long)y);         y = pos[d];       }     }     return -1;   }   int main() {     while (scanf( %d%d%*c , &n, &m) != EOF) {       for (int i = 0; i < n; i++) gets(mat[i]);       long long ans = solve();       if (ans == -1)         puts( Never );       else         cout << ans << endl;     }     return 0;   }
#include <bits/stdc++.h>   using namespace std;   int t[1000005], r = 1000004;   int dp[1000005];   int main() {     memset(dp, 0, sizeof(dp));     memset(t, 0, sizeof(t));     int n, ans = 0, x;     cin >> n;     for (int i = 0; i < n; i++) {       cin >> x;       cin >> t[x];     }     dp[0] = 0;     if (t[0] != 0) dp[0] = 1;     for (int i = 1; i <= r; i++) {       if (!t[i])         dp[i] = dp[i - 1];       else {         if (t[i] >= i)           dp[i] = 1;         else           dp[i] = dp[i - t[i] - 1] + 1;       }       ans = max(ans, dp[i]);     }     cout << n - ans;     return 0;   }
#include <bits/stdc++.h>   using namespace std;   using ll = long long;   const int N = 2e5 + 5;   int n, q;   ll a[N], tree[N << 2], lz[N << 2];   void push(int node, int s, int e) {     tree[node] += lz[node];     if (s != e) {       lz[node << 1] += lz[node];       lz[node << 1 | 1] += lz[node];     }     lz[node] = 0;   }   void upd(int node, int s, int e, int l, int r, ll v) {     push(node, s, e);     if (r < s || e < l || l > r) return;     if (l <= s && e <= r) {       lz[node] += v;       push(node, s, e);       return;     }     int md = (s + e) >> 1;     upd(node << 1, s, md, l, r, v);     upd(node << 1 | 1, md + 1, e, l, r, v);     tree[node] = max(tree[node << 1], tree[node << 1 | 1]);   }   int get(int node, int s, int e) {     push(node, s, e);     if (tree[node] < 0) return -1;     if (s == e) return tree[node] != 0 ? -1 : s + 1;     int md = (s + e) >> 1;     int x = get(node << 1, s, md);     if (~x) return x;     return get(node << 1 | 1, md + 1, e);   }   int main() {     scanf( %d%d , &n, &q);     for (int i = 0; i < n; i++) {       scanf( %lld , a + i);       upd(1, 0, n - 1, i, i, a[i]);       upd(1, 0, n - 1, i + 1, n - 1, -a[i]);     }     while (q--) {       int x, v;       scanf( %d%d , &x, &v);       upd(1, 0, n - 1, x - 1, x - 1, -a[x - 1] + v);       upd(1, 0, n - 1, x, n - 1, a[x - 1] - v);       a[x - 1] = v;       printf( %d n , get(1, 0, n - 1));     }     return 0;   }
#include <bits/stdc++.h>   const double pi = acos(-1.0), eps = 1e-9;   const int dx[8] = {1, -1, 0, 0, 1, 1, -1, -1};   const int dy[8] = {0, 0, 1, -1, 1, -1, -1, 1};   const int MO = (int)(1e9 + 7);   using namespace std;   int main() {     long long n;     cin >> n;     while (n % 3 == 0) n /= 3;     cout << n / 3 + 1 << endl;   }
#include <bits/stdc++.h>   using namespace std;   void fileio() {}   vector<vector<long long int>> g;   long long int cube_root(long long int x) {     long long int lo = 0, hi = 1e6 + 10;     while (lo <= hi) {       long long int m = (lo + hi) / 2;       long long int val = m * m * m;       if (val == x) return m;       if (val > x) {         hi = m - 1;       } else {         lo = m + 1;       }     }     return -1;   }   int32_t main() {     ios_base::sync_with_stdio(false);     cin.tie(NULL);     fileio();     long long int t;     cin >> t;     while (t-- > 0) {       long long int a, b;       cin >> a >> b;       long long int x = a * b;       long long int c = cube_root(x);       if (c == -1) {         cout <<  NO ;         cout <<   n ;         continue;       }       if (a % c != 0 or b % c != 0) {         cout <<  NO ;         cout <<   n ;         continue;       }       cout <<  YES ;       cout <<   n ;     }     return 0;   }
#include <bits/stdc++.h>   using namespace std;   int vis[51];   int main() {     string s;     int n, len;     cin >> len >> n;     cin >> s;     while (n--) {       for (int i = 0; i <= len; i++) {         if (s[i] ==  B  && s[i + 1] ==  G ) {           swap(s[i], s[i + 1]);           i++;         }       }     }     cout << s << endl;     return 0;   }
#include <bits/stdc++.h>   using namespace std;   int main() {     int N;     cin >> N;     int a[10000], b[10000];     for (int n = 0; n < N; n++) {       cout <<  ?   << 0 <<     << n << endl;       cin >> a[n];     }     for (int n = 0; n < N; n++) {       cout <<  ?   << n <<     << 0 << endl;       cin >> b[n];     }     int nm = 0, fa = 0;     int ac[10000], bc[10000];     for (int n = 0; n < N; n++) {       for (int x = 0; x < N; x++) {         ac[x] = a[x] ^ n;       }       for (int x = 0; x < N; x++) {         bc[x] = b[x] ^ ac[0];       }       bool p = true;       for (int x = 0; x < N; x++) {         if (ac[bc[x]] != x) {           p = false;         }       }       if (p) {         fa = n;         nm++;       }     }     cout <<  !  << endl;     cout << nm << endl;     fa = fa ^ a[0];     for (int x = 0; x < N; x++) {       if (x) cout <<    ;       cout << (fa ^ b[x]);     }     cout << endl;     return 0;   }
#include <bits/stdc++.h>   using namespace std;   int ans1 = 0, ans2 = 0;   int main(void) {     int t, i, j;     cin >> t;     while (t--) {       int n, mx = INT_MIN, ans = 0, sz;       cin >> n;       int a[n];       set<int> s;       for (i = 0; i < n; i++) {         cin >> a[i];         mx = max(mx, a[i]);         s.insert(a[i]);       }       for (i = 12; i >= 0; i--) {         if (((1 << i) & mx) != 0) break;       }       mx = (1 << (i + 1));       for (i = 1; i <= mx; i++) {         map<int, int> mp;         sz = 0;         for (j = 0; j < n; j++) {           if ((s.count(a[j] ^ i) > 0) && (mp[a[j] ^ i] == 0)) {             sz++;             mp[a[j] ^ i]++;           }         }         if (sz == s.size()) ans = 1;         if (ans == 1) break;       }       if (ans == 0)         cout << -1 << endl;       else         cout << i << endl;     }   }
#include <bits/stdc++.h>   using namespace std;   int main() {     int n, k, a, b;     scanf( %d%d , &n, &k);     printf( %d , n);     for (int i = n - 1; i > k; --i) printf(  %d , i);     a = 1, b = k;     while (a < b) printf(  %d %d , a++, b--);     if (a == b) printf(  %d , a);     puts(  );   }
#include <bits/stdc++.h>   using namespace std;   int main() {     ios_base::sync_with_stdio(false);     cin.tie(NULL);     int n, m, c;     cin >> n >> m >> c;     int a[n], b[m];     for (int i = 0; i < n; i++) cin >> a[i];     for (int i = 0; i < m; i++) cin >> b[i];     for (int i = 0; i < n - m + 1; i++) {       int k = i;       for (int j = 0; j < m; j++) {         a[k] += b[j];         a[k] = a[k] % c;         k++;       }     }     for (int i = 0; i < n; i++) cout << a[i] <<    ;   }
#include <bits/stdc++.h>   using namespace std;   long long a, b, c, l, tl, ans, ca, cb, cc;   long long min(long long a, long long b) { return a < b ? a : b; }   void cpy() {     a = ca, b = cb;     c = cc;   }   int main() {     scanf( %lld%lld%lld%lld , &a, &b, &c, &l);     ca = a;     cb = b;     cc = c;     for (int i = l + 2; i >= 2; --i) ans += 1ll * i * (1ll * i - 1) / 2;     if (b + c - a <= l) {       if (a > b + c)         tl = l;       else         tl = l - (b + c - a), a = b + c;       for (int i = 0; i <= tl; ++i) {         long long tb = min(a - b - c + 1ll * i, tl - 1ll * i);         ans -= (tb + 1) * (tb + 2) / 2;       }     }     cpy();     if (a + b - c <= l) {       if (c > a + b)         tl = l;       else         tl = l - (a + b - c), c = a + b;       for (int i = 0; i <= tl; ++i) {         long long tb = min(c - a - b + i * 1ll, tl - i * 1ll);         ans -= (tb + 1) * (tb + 2) / 2;       }     }     cpy();     if (a + c - b <= l) {       if (b > a + c)         tl = l;       else         tl = l - (a + c - b), b = a + c;       for (int i = 0; i <= tl; ++i) {         long long tb = min(b - a - c + i * 1ll, tl - i * 1ll);         ans -= (tb + 1) * (tb + 2) / 2;       }     }     printf( %lld , ans);     return 0;   }
#include <bits/stdc++.h>   using namespace std;   string a, b;   bool func(string s, long n) {     if (s.size() % n) return false;     for (long i = 0; i < s.size(); i++)       if (s[i] != a[i % n]) return 0;     return 1;   }   int main() {     cin >> a >> b;     long ans = 0, i = 1;     string p = a, q = b;     long minx = min(a.size(), b.size());     while (i <= minx) {       ans += (func(p, i) && func(q, i));       i++;     }     cout << ans;   }
#include <bits/stdc++.h>   using namespace std;   #pragma warning(disable : 4996)   const double PI = acos(-1);   const long long N = 500005, md = 1e9 + 7, inf = 1e18, iinf = 1e9 + 5, MX = 3e5;   long long g, w, h, x, y, z, i, j, k, q, n, cnt, res, sum, r[N], a[N], b[N];   char c[N];   string s, t;   pair<int, int> p[N];   vector<int> e[N];   int main() {     ios::sync_with_stdio(false), cin.tie(0);     cin >> n >> k;     if (k == 2 && (n < 3 || n == 4))       cout <<  NO ;     else if (k == 3 && (n < 6 || n == 8))       cout <<  NO ;     else {       if (n < k * (k + 1) / 2) cout <<  NO , exit(0);       cout <<  YES n ;       x = (n - k * (k + 1) / 2) / k;       for (long long i = 0; i < k; i++) a[i] = i + 1 + x, sum += a[i];       while (sum < n) {         for (long long i = k - 1; i >= 0; i--) {           if (a[i] + 1 > a[i - 1] * 2) i = k - 1;           if (sum == n) break;           if (sum < n) sum++, a[i]++;         }       }       for (long long i = 0; i < k; i++) cout << a[i] <<    ;     }   }
#include <bits/stdc++.h>   using namespace std;   bool ask(int val) {     cout <<  ?   << val + 1 << endl;     string resp;     cin >> resp;     return resp[0] ==  Y  ? true : false;   }   void fix() { cout <<  R  << endl; }   int main() {     int n, k;     cin >> n >> k;     vector<bool> check(n + 1, 0);     int block = k % 2 ? 1 : k / 2;     for (int blockno = 1; blockno < n / block; blockno += 1) {       for (int j = 0; j < blockno && (blockno + j < n / block); j += 1) {         for (int k = j; k < n / block; k += blockno) {           for (int a = 0; a < block; a += 1) {             if (ask(k * block + a)) {               check[k * block + a + 1] = true;             }           }         }         fix();       }     }     int res = 0;     for (int i = 1; i <= n; i += 1) {       if (check[i] == 0) {         res += 1;       }     }     cout <<  !   << res << endl;   }
#include <bits/stdc++.h>   using namespace std;   int _ = 1, tc = 1;   int dx[] = {0, 0, 1, -1};   int dy[] = {1, -1, 0, 0};   const int MAXN = 1e7 + 1;   int spf[MAXN];   int d, x, y;   void sieve() {     for (int i = 1; i < MAXN; i++) spf[i] = i;     for (int i = 4; i < MAXN; i += 2) spf[i] = 2;     for (int i = 3; i * i < MAXN; i += 2) {       if (spf[i] == i) {         for (int j = i * i; j < MAXN; j += i) {           if (spf[j] == j) spf[j] = i;         }       }     }   }   void extended_gcd(int a, int b) {     if (b == 0) {       d = a;       x = 1;       y = 0;     } else {       extended_gcd(b, a % b);       int temp = y;       y = x - y * (a / b);       x = temp;     }   }   int gcd(int a, int b) { return b == 0 ? a : gcd(b, a % b); }   int lcm(int a, int b) { return a * (b / gcd(a, b)); }   int getmod(int x, int y, int mod) { return (((x - y) % mod) + mod) % mod; }   vector<string> split(string line, char dil) {     vector<string> tokens;     stringstream check1(line);     string intermediate;     while (getline(check1, intermediate, dil)) {       tokens.push_back(intermediate);     }     return tokens;   }   void read(vector<long long>& v) {     for (int i = 0; i < v.size(); i++) cin >> v[i];   }   long long calc(long long curr, long long cuts) {     int p2 = curr % cuts;     int p1 = cuts - p2;     long long w = curr / cuts;     long long val = p1 * w * w + p2 * (w + 1) * (w + 1);     return val;   }   void solve() {     int n, k;     cin >> n >> k;     vector<long long> v(n);     read(v);     priority_queue<pair<long long, pair<long long, long long> > > q;     long long time = 0;     for (int i = 0; i < n; i++) {       long long curr = v[i];       time += curr * curr;       q.push(make_pair(calc(curr, 1) - calc(curr, 2), make_pair(curr, 2)));     }     while (n < k) {       pair<long long, pair<long long, long long> > curr = q.top();       q.pop();       long long dec = curr.first;       long long past = curr.second.second;       long long org = curr.second.first;       q.push(make_pair(calc(org, past) - calc(org, past + 1),                        make_pair(org, past + 1)));       n++;       time -= dec;     }     cout << time << endl;   }   int main() {     ios_base::sync_with_stdio(false);     cin.tie(0);     cout.tie(0);     while (_--) {       solve();     }     return 0;   }
#include <bits/stdc++.h>   using namespace std;   long long gcd(long long a, long long b) { return b == 0 ? a : gcd(b, a % b); }   const int MAXN = 1000;   int n;   char s[MAXN + 1];   int solve() {     int ret = n;     for (int i = (0); i < (n); ++i)       if (s[i] ==  1  || s[n - i - 1] ==  1 ) ret = max(ret, 2 * (n - i));     return ret;   }   void run() {     scanf( %d , &n);     scanf( %s , s);     assert(strlen(s) == n);     printf( %d n , solve());   }   int main() {     int ncase;     scanf( %d , &ncase);     for (int i = (1); i <= (ncase); ++i) run();     return 0;   }