重构 countBits 函数，修复返回值并添加 ASSERT 宏以增强测试

src/338.c

@@ -1,13 +1,44 @@
 //
 // Created by aurora on 2024/9/16.
 //
+#include <assert.h>
 #include <solution/338.h>
+#include <stdlib.h>
+
+#define ASSERT(x, expect)                                                                                              \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        int retSize;                                                                                                   \
+        int *ret;                                                                                                      \
+        ret = countBits(x, &retSize);                                                                                  \
+        for (int i = 0; i < retSize; i++)                                                                              \
+            assert(ret[i] == expect[i]);                                                                               \
+        free(ret);                                                                                                     \
+    } while (0)
+
 int *countBits(int n, int *returnSize)
 {
-    return 0;
+    *returnSize = n + 1;
+    int *ret = (int *)malloc(sizeof(int) * (*returnSize));
+    ret[0] = 0;
+    for (int i = 1; i <= n; i++)
+    {
+        ret[i] = ret[i >> 1] + (i & 1);
+        // equal to the following expression
+        // if (i % 2 == 0)
+        //     ret[i] = ret[i / 2];
+        // else
+        //     ret[i] = ret[i - 1] + 1;
+    }
+
+    return ret;
 }
 
-int main() {
+int main()
-
+{
+    int t1[] = {0, 1, 1};
+    ASSERT(2, t1);
+    int t2[] = {0, 1, 1, 2, 1, 2};
+    ASSERT(5, t2);
     return 0;
 }