[Hackerearth] Special Subarrays

 

Problem Link    : Special Subarrays
Category        : Trie
Contest         : DSA Coding Contest - January '24

#include "bits/stdc++.h"
using namespace std;
#define int     long long int
#define endl    '\n'
 
template <const int32_t MOD> struct modint {
    int32_t value;
    modint() = default;
    modint(int32_t value_) : value(value_) {}
    inline modint<MOD> operator + (modint<MOD> other) const { int32_t c = this->value + other.value; return modint<MOD>(c >= MOD ? c - MOD : c); }
    inline modint<MOD> operator - (modint<MOD> other) const { int32_t c = this->value - other.value; return modint<MOD>(c <    0 ? c + MOD : c); }
    inline modint<MOD> operator * (modint<MOD> other) const { int32_t c = (int64_t)this->value * other.value % MOD; return modint<MOD>(c < 0 ? c + MOD : c); }
    inline modint<MOD> & operator += (modint<MOD> other) { this->value += other.value; if (this->value >= MOD) this->value -= MOD; return *this; }
    inline modint<MOD> & operator -= (modint<MOD> other) { this->value -= other.value; if (this->value < 0) this->value += MOD; return *this; }
    inline modint<MOD> & operator *= (modint<MOD> other) { this->value = (int64_t)this->value * other.value % MOD; if (this->value < 0) this->value += MOD; return *this; }
    inline modint<MOD> operator - () const { return modint<MOD>(this->value ? MOD - this->value : 0); }
    modint<MOD> pow(uint64_t k) const { modint<MOD> x = *this, y = 1; for (; k; k >>= 1) { if (k & 1) y *= x; x *= x; } return y; }
    modint<MOD> inv() const { return pow(MOD - 2); }  // MOD must be a prime
    inline modint<MOD> operator /  (modint<MOD> other) const { return *this *  other.inv(); }
    inline modint<MOD> operator /= (modint<MOD> other)       { return *this *= other.inv(); }
    inline bool operator == (modint<MOD> other) const { return value == other.value; }
    inline bool operator != (modint<MOD> other) const { return value != other.value; }
    inline bool operator < (modint<MOD> other) const { return value < other.value; }
    inline bool operator > (modint<MOD> other) const { return value > other.value; }
};
template <int32_t MOD> modint<MOD> operator * (int64_t value, modint<MOD> n) { return modint<MOD>(value) * n; }
template <int32_t MOD> modint<MOD> operator * (int32_t value, modint<MOD> n) { return modint<MOD>(value % MOD) * n; }
template <int32_t MOD> istream & operator >> (istream & in, modint<MOD> &n) { return in >> n.value; }
template <int32_t MOD> ostream & operator << (ostream & out, modint<MOD> n) { return out << n.value; }
 
const int mod = 1e9 + 7;
using mint = modint<mod>;
 
struct Trie {
    int ones;
    Trie *child[2];
 
    Trie() {
        ones = 0;
        for (int i = 0; i < 2; i++) {
            child[i] = nullptr;
        }
    }
    ~Trie() {
        for (int i = 0; i < 2; i++) {
            if (child[i] and this != child[i]) {
                delete child[i];
            }
        }
    }
};
 
typedef Trie* pnode;
pnode root;
 
void add(string &s) {
    pnode curRoot = root;
    for (char ch : s) {
        int val = ch - '0';
        if (curRoot->child[val] == nullptr) {
            curRoot->child[val] = new Trie();
        }
        curRoot = curRoot->child[val];
    }
}
 
mint ans;
 
void dfs(pnode curRoot, int ones) {
    for (int i = 0; i < 2; i++) {
        if (curRoot->child[i]) {
            int nones = ones + (i == 0);
            ans += max(0ll, nones - 1);
            dfs(curRoot->child[i], nones);
        }
    }
}
 
signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);
    cout.precision(12);
 
    bool FILEIO = 1;
    if (FILEIO and fopen("in.txt", "r")) {
        freopen("in.txt", "r", stdin);
    }
 
    int n;
    cin >> n;
    root = new Trie();
    for (int i = 0; i < n; i++) {
        string s;
        cin >> s;
        add(s);
    }
    ans = 0;
    dfs(root, 0);
    cout << ans << endl;
}