select for bit arrays

data/BitArray.cpp

@@ -47,10 +47,17 @@ static int getArrayLength(int length, int bits) {
 BitArray::BitArray() : length(0), bits(0), data(nullptr) {
 }
 
+static int* allocate(int length, int bits) {
+    int l = getArrayLength(length, bits);
+    int* a = new int[l];
+    memset(a, 0, l * sizeof(int));
+    return a;
+}
+
 BitArray::BitArray(int length, int bits)
     : length(length), bits(bits), data(nullptr) {
     if(length > 0 && bits > 0) {
-        data = new int[getArrayLength(length, bits)];
+        data = allocate(length, bits);
     }
 }
 
@@ -97,6 +104,30 @@ int BitArray::getInternalByteSize() const {
     return getArrayLength(length, bits) * sizeof(int);
 }
 
+int BitArray::select(int index) const {
+    if(index <= 0) {
+        return -1;
+    }
+    int found = 0;
+    int end = getArrayLength(length, bits);
+    for(int i = 0; i < end; i++) {
+        int ones = __builtin_popcount(data[i]);
+        found += ones;
+        if(found >= index) {
+            found -= ones;
+            int a = i * 32 - 1;
+            int d = data[i];
+            while(found < index) {
+                found += d & 1;
+                d >>= 1;
+                a++;
+            }
+            return a;
+        }
+    }
+    return -1;
+}
+
 void BitArray::fill(int value) {
     for(int i = 0; i < length; i++) {
         set(i, value);
@@ -104,7 +135,7 @@ void BitArray::fill(int value) {
 }
 
 void BitArray::resize(int newLength, int newBits) {
-    int* newData = new int[getArrayLength(newLength, newBits)];
+    int* newData = allocate(newLength, newBits);
     int end = Math::min(length, newLength);
     for(int i = 0; i < end; i++) {
         setBits(newData, i, newBits, get(i));

data/BitArray.h

@@ -23,6 +23,8 @@ public:
     int getBits() const;
     int getInternalByteSize() const;
 
+    int select(int index) const;
+
     void resize(int newLength, int newBits);
 
     void fill(int value);

tests/BitArrayTests.cpp

@@ -61,6 +61,23 @@ static void testReadOnly(Test& test) {
     test.checkEqual(4, bits2.get(3), "can read from const 4");
 }
 
+static void testSelect(Test& test) {
+    BitArray bits(90, 1);
+    bits.fill(0);
+    bits.set(0, 1).set(5, 1).set(20, 1).set(31, 1);
+    bits.set(32, 1).set(33, 1).set(60, 1);
+    test.checkEqual(-1, bits.select(-1), "select -1 bit");
+    test.checkEqual(-1, bits.select(0), "select 0 bit");
+    test.checkEqual(0, bits.select(1), "select 1 bit");
+    test.checkEqual(5, bits.select(2), "select 2 bit");
+    test.checkEqual(20, bits.select(3), "select 3 bit");
+    test.checkEqual(31, bits.select(4), "select 4 bit");
+    test.checkEqual(32, bits.select(5), "select 5 bit");
+    test.checkEqual(33, bits.select(6), "select 6 bit");
+    test.checkEqual(60, bits.select(7), "select 7 bit");
+    test.checkEqual(-1, bits.select(8), "select 8 bit");
+}
+
 static void testToString1(Test& test) {
     BitArray bits(4, 3);
     bits.set(0, 1).set(1, 2).set(2, 3).set(3, 4);
@@ -80,6 +97,7 @@ void BitArrayTests::test() {
     testBigSetRead(test);
     testRandomSetReadResize(test);
     testReadOnly(test);
+    testSelect(test);
     testToString1(test);
     testToString2(test);
     test.finalize();