aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--count_pairs_xor_range.c88
-rw-r--r--count_pairs_xor_range.py35
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))