Google Kickstart round G 2018（状压搜索）

A

输出满足条件的三元组的个数

特殊考虑0和1

WA2

细节写的太丑了

#include<iostream>
#include<string.h>
#include<algorithm>
#define N 10005
using namespace std;
long long a[N];
long long sum[200005];
int main(){
    int T,ca=0;
    scanf("%d",&T);
    while (T--){
        ca++;
        memset(sum,0,sizeof(sum));
        int n;
        scanf("%d",&n);
        for (int i=1;i<=n;i++){
            scanf("%lld",&a[i]);
            sum[a[i]]++;
        }
        long long  ans=0;
        sort(a+1,a+1+n);
        long long tot1=0,tot0=0;
        a[n+1]=2;
        for (int i=1;i<=n+1;i++) if (a[i]!=0){ tot0=i-1;break; }
        if (tot0>=3) ans+=(tot0*(tot0-1)*(tot0-2)/6);
        for (int i=tot0+1;i<=n+1;i++) if (a[i]!=1){ tot1=i-1-tot0;