นับ Subarray ที่มีจำนวน 0 และ 1 เท่ากันใน C++

เราได้รับอาร์เรย์ arr[] ที่มีแต่ 0 และ 1 เป้าหมายคือนับ subarray ที่มีจำนวน 0 และ 1 เท่ากัน

ตัวอย่าง

อินพุต: arr[] = {0, 0, 1, 1, 1, 0}

เอาต์พุต: 4

คำอธิบาย: [0,0,1,1], [0,1], [1,0], [0,0,1,1,1,0]

อินพุต: arr[] = {0, 1, 1, 1, 1}

เอาต์พุต: 1

คำอธิบาย: [0,1] เท่านั้น

แนวทางการแก้ปัญหา (Naive O(n²))

ใช้ nested loop สร้าง subarray ทุกคู่ นับ 0 และ 1 ในแต่ละ subarray ถ้าเท่ากันนับเพิ่ม

การนำไปใช้ (Implementation)

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

int sub_zeroes_ones(int arr[], int size) {
    int count = 0;
    for (int i = 0; i <= size - 1; i++) {
        int total_0 = 0, total_1 = 0;
        for (int j = i; j <= size - 1; j++) {
            if (arr[j] == 0) total_0++;
            else if (arr[j] == 1) total_1++;
            if (total_0 == total_1) count++;
        }
    }
    return count;
}

int main() {
    int arr[] = {0, 1, 1, 0, 0};
    int size = sizeof(arr) / sizeof(arr[0]);
    cout << "Count of subarrays with equal number of 1's and 0's are: " << sub_zeroes_ones(arr, size);
    return 0;
}

เอาต์พุต

Count of subarrays with equal number of 1's and 0's are: 4

หมายเหตุ: วิธีนี้เป็น O(n²) สำหรับข้อมูลใหญ่ควรใช้ prefix sum + hash map เพื่อลดเป็น O(n)