diff options
-rw-r--r-- | count_pairs_xor_range.c | 88 | ||||
-rw-r--r-- | count_pairs_xor_range.py | 35 |
2 files changed, 123 insertions, 0 deletions
diff --git a/count_pairs_xor_range.c b/count_pairs_xor_range.c new file mode 100644 index 0000000..7f0b07f --- /dev/null +++ b/count_pairs_xor_range.c @@ -0,0 +1,88 @@ +// 1803. Count Pairs With XOR in a Range +#include "leetcode.h" + +/* + * given a 0-indexed integer array 'numss' and two integers 'low' and 'high', + * return the numsber of nice pairs. a nice pair is a pair '(i, j)' where '0 <= + * i < j < numss.len' and 'low <= (numss[i] ^ numss[j]) <= high'. + */ + +#pragma GCC optimize("O3") +#pragma GCC target("avx2") +#include <x86intrin.h> + +/* + * the following code is heavily optimised with avx2 intrinsics and has zero + * branching. the general structure follows an intuitive brute force approach: + * + * int pairs = 0; + * for (int i = 0; i < numsSize - 1; i++) + * for (int j = i + 1; j < n; j++) + * if ((nums[i] ^ nums[j]) >= low && (nums[i] ^ nums[j]) <= high) + * pairs++; + * return pairs; + */ + +int countPairs(int *nums, int numsSize, int low, int high) { + int pairs = 0; + + /* + * store input as 16-bit integers, allowing 16 pairs to be processed per + * iteration. 32 bytes of extra entries are added to avoid kludge for last + * entries. the inner loop will process the array in chunks of 16 shorts per + * execution + */ + unsigned short __attribute__((aligned(32))) d[numsSize + 16]; + for (int i = 0; i < numsSize; i++) + d[i] = nums[i]; + for (int i = numsSize; i < numsSize + 16; i++) + d[i] = 0xFFFF; + + /* + * create arrays for lower and upper bound. the '_m256i' type is a 256-bit + * register that can hold 16 shorts at a time. two registers are allocated + * with each 16-bit value set to low and high respectively + */ + const __m256i low_mask = _mm256_set1_epi16((short)low); + const __m256i high_mask = _mm256_set1_epi16((short)high); + + for (int i = 0; i < numsSize - 1; i++) { + // iterate through elements and create a register array 'a[i]' where 16 + // elements are set to 'd[i]' + __m256i ai = _mm256_set1_epi16(d[i]); + + // critical section + for (int j = i + 1; j < numsSize; j += 16) { + // read 16 numbers to array 'a[j]' + __m256i aj = _mm256_lddqu_si256((__m256i *)(d + j)); + + // XOR them with array 'a[i]', essentially creating array of '{d[j] ^ + // d[i], d[j + 1] ^ d[i], ..., d[j + 15] ^ d[i]}' + __m256i xij = _mm256_xor_si256(ai, aj); + + /* + * compare XOR results with upper and lower bounds. the 'cmpgt' intrinsic + * creates an array of compare results, where an element is set to + * '0xFFFF' when the corresponding element in 'xij' is greater than the + * corresponding element in 'high_mask' + */ + __m256i high_ans = _mm256_cmpgt_epi16(xij, high_mask); + __m256i low_ans = _mm256_cmpgt_epi16(low_mask, xij); + + /* + * OR the comparison results to mask, count the number of high bits set, + * and add it to 'pairs'. note, the comparison sets bits when the XOR is + * outside the '[low, high]' interval + */ + __m256i ans = _mm256_or_si256(high_ans, low_ans); + pairs += 32 - __builtin_popcount(_mm256_movemask_epi8(ans)); + } + } + return pairs >> 1; +} + +int main() { + int n1[] = {1, 4, 2, 7}, n2[] = {9, 8, 4, 2, 1}; + printf("%d\n", countPairs(n1, ARRAY_SIZE(n1), 2, 6)); // expect: 6 + printf("%d\n", countPairs(n2, ARRAY_SIZE(n2), 5, 14)); // expect: 8 +} diff --git a/count_pairs_xor_range.py b/count_pairs_xor_range.py new file mode 100644 index 0000000..01f2b84 --- /dev/null +++ b/count_pairs_xor_range.py @@ -0,0 +1,35 @@ +# 1803. Count Pairs With XOR in a Range +from collections import Counter + +""" +given a 0-indexed integer array 'numss' and two integers 'low' and 'high', +return the numsber of nice pairs. a nice pair is a pair '(i, j)' where '0 <= +i < j < numss.len' and 'low <= (numss[i] ^ numss[j]) <= high'. +""" + + +class Solution(object): + def countPairs(self, nums, low, high): + """ + :type nums: List[int] + :type low: int + :type high: int + :rtype: int + """ + + def check(nums, x): + cnt, res = Counter(nums), 0 + while x: + if x & 1: + res += sum(cnt[i] * cnt[(x - 1) ^ i] for i in cnt) + cnt = Counter({i >> 1: cnt[i] + cnt[i + 1] for i in cnt}) + x >>= 1 + return res // 2 + + return check(nums, high + 1) - check(nums, low) + + +if __name__ == "__main__": + obj = Solution() + print(obj.countPairs(nums=[1, 4, 2, 7], low=2, high=6)) + print(obj.countPairs(nums=[9, 8, 4, 2, 1], low=5, high=14)) |