#include <iostream>
using namespace std;
int main() {
// your code goes here
int test;
cin>>test;
while(test--)
{
int n;
cin>>n;
char a[n];
int i, count=0, ans=0;
for(i=0; i<n; i++)
{
cin>>a[i];
if(a[i]=='1')
count++;
}
cout<<(count*(count+1)/2)<<endl;
}
return 0;
}
Below is the code that works
#include <bits/stdc++.h>
using namespace std;
int main()
{
int t;
cin>>t;
while(t--) {
int n;
cin>>n;
char a[n];
int i, ans=0;
long long int count=0;
for(i=0; i<n; i++)
{
cin>>a[i];
if(a[i]=='1')
count++;
}
cout<<(count*(count+1)/2)<<endl;
}
return 0;
}
I have only changed the datatype of count variable from int to long long int.
In your solution the datatype of count is int so when you calculate count*(count+1)/2 then the overall value would be returned as int.
Now, consider a case when length of string is 10^5 and all characters are 1. In that case n*(n+1)/2 would exceed the limit of int and hence give wrong answer.
Thats why long long int will be used as the datatype of count variable here.