8namespace py = pybind11;
10template <typename T> py::array get_run_lengths_impl(py::array_t<T> x) {
11 auto x_buf = x.request();
12 T *x_ptr = static_cast<T *>(x_buf.ptr);
13 size_t N = x_buf.shape[0];
15 std::vector<uint32_t> runs;
16 size_t i = 0;
17 uint32_t current_nonzero_run_length = 0;
19 while (i < N) {
20 // Check for a sequence of at least 10 zeros
21 bool has_zeros = true;
22 for (size_t j = 0; j < 10 && i + j < N; j++) {
23 if (x_ptr[i + j] != 0) {
24 has_zeros = false;
25 break;
26 }
27 }
29 if (has_zeros) {
30 // Add current non-zero run if any
31 runs.push_back(current_nonzero_run_length);
32 current_nonzero_run_length = 0;
34 // Count consecutive zeros
35 size_t j = i;
36 while (j < N && x_ptr[j] == 0) {
37 j++;
38 }
39 runs.push_back(j - i);
40 i = j;
41 } else {
42 current_nonzero_run_length++;
43 i++;
44 }
45 }
47 // Add final non-zero run if any
48 if (current_nonzero_run_length > 0) {
49 runs.push_back(current_nonzero_run_length);
50 }
52 // Determine appropriate dtype based on max run length
53 uint32_t max_run = 0;
54 for (const auto &run : runs) {
55 if (run > max_run) {
56 max_run = run;
57 }
58 }
60 // Create numpy array with appropriate dtype
61 std::vector<ssize_t> shape = {static_cast<ssize_t>(runs.size())};
63 if (max_run < 256) {
64 py::array_t<uint8_t> result(shape);
65 auto result_buf = result.request();
66 uint8_t *result_ptr = static_cast<uint8_t *>(result_buf.ptr);
67 for (size_t i = 0; i < runs.size(); i++) {
68 result_ptr[i] = static_cast<uint8_t>(runs[i]);
69 }
70 return result;
71 } else if (max_run < 65536) {
72 py::array_t<uint16_t> result(shape);
73 auto result_buf = result.request();
74 uint16_t *result_ptr = static_cast<uint16_t *>(result_buf.ptr);
75 for (size_t i = 0; i < runs.size(); i++) {
76 result_ptr[i] = static_cast<uint16_t>(runs[i]);
77 }
78 return result;
79 } else {
80 py::array_t<uint32_t> result(shape);
81 auto result_buf = result.request();
82 uint32_t *result_ptr = static_cast<uint32_t *>(result_buf.ptr);
83 for (size_t i = 0; i < runs.size(); i++) {
84 result_ptr[i] = runs[i];
85 }
86 return result;
87 }
88}