python - 从两个数组中获取相同数量的元素，使得所取值的重复项尽可能少

Question

假设我们有 2 个大小为 N 的数组，它们的值在 [0, N-1] 范围内。例如:

a = np.array([0, 1, 2, 0])
b = np.array([2, 0, 3, 3])

我需要生成一个新数组 c，它恰好包含来自 a 和 b 的 N/2 个元素分别，即必须从两个父数组中均匀/相等地获取值。

(对于奇数长度，这将是 (N-1)/2 和 (N+1)/2。也可以忽略奇数长度的情况，不重要).

从两个数组中获取相同数量的元素非常简单，但还有一个额外的限制:c 应该有尽可能多的唯一数字/尽可能少的重复项。

比如上面的a和b的解决方案是:

c = np.array([b[0], a[1], b[2], a[3]])
>>> c
array([2, 1, 3, 0])

请注意，位置/顺序已保留。 a 和 b 的每个元素，我们用来形成 c 的位置相同。如果 c 中的元素 i 来自 a，则 c[i] == a[i]，同理b.

---

一个简单的解决方案就是一种简单的路径遍历，很容易递归实现:

def traverse(i, a, b, path, n_a, n_b, best, best_path):
    if n_a == 0 and n_b == 0:
        score = len(set(path))
        return (score, path.copy()) if score > best else (best, best_path)

    if n_a > 0:
        path.append(a[i])
        best, best_path = traverse(i + 1, a, b, path, n_a - 1, n_b, best, best_path)
        path.pop()
    
    if n_b > 0:
        path.append(b[i])
        best, best_path = traverse(i + 1, a, b, path, n_a, n_b - 1, best, best_path)
        path.pop()

    return best, best_path

这里的n_a和n_b分别是我们要从a和b中取多少个值，它是2 和 2 因为我们想平均取 4 项。

>>> score, best_path = traverse(0, a, b, [], 2, 2, 0, None)
>>> score, best_path
(4, [2, 1, 3, 0])

---

是否有可能通过 numpy 以更向量化/更高效的方式实现上述内容？

Answer 1

该算法之所以慢主要是因为它在指数级 时间内运行。由于递归，没有直接使用 Numpy 对该算法进行矢量化的直接方法。即使有可能，大量的组合也会导致大多数 Numpy 实现效率低下(由于要计算的 Numpy 数组很大)。此外，AFAIK 没有矢量化操作来有效地计算许多行的唯一值的数量(通常的方法是使用 np.unique 在这种情况下效率不高，如果没有循环就不能使用)。因此，有两种可能的策略来加快速度:

• 尝试找到具有合理复杂度的算法(例如 <= O(n^4) )；
• 使用编译方法、微优化和技巧来编写更快的暴力实现。

由于找到正确的次指数算法并不容易，我选择了另一种方法(尽管第一种方法是最好的)。

想法是:

• 通过使用整数循环迭代生成所有可能的解决方案来消除递归；
• 编写一种快速计算数组中唯一项的方法；
• 使用 Numba JIT 编译器来优化只有在编译后才有效的代码。

这是最终代码:

import numpy as np
import numba as nb

# Naive way to count unique items.
# This is a slow fallback implementation.
@nb.njit
def naive_count_unique(arr):
    count = 0
    for i in range(len(arr)):
        val = arr[i]
        found = False
        for j in range(i):
            if arr[j] == val:
                found = True
                break
        if not found:
            count += 1
    return count

# Optimized way to count unique items on small arrays.
# Count items 2 by 2.
# Fast on small arrays.
@nb.njit
def optim_count_unique(arr):
    count = 0
    for i in range(0, len(arr), 2):
        if arr[i] == arr[i+1]:
            tmp = 1
            for j in range(i):
                if arr[j] == arr[i]: tmp = 0
            count += tmp
        else:
            val1, val2 = arr[i], arr[i+1]
            tmp1, tmp2 = 1, 1
            for j in range(i):
                val = arr[j]
                if val == val1: tmp1 = 0
                if val == val2: tmp2 = 0
            count += tmp1 + tmp2
    return count

@nb.njit
def count_unique(arr):
    if len(arr) % 2 == 0:
        return optim_count_unique(arr)
    else:
        # Odd case: not optimized yet
        return naive_count_unique(arr)

# Count the number of bits in a 32-bit integer
# See https://stackoverflow.com/questions/71097470/msb-lsb-popcount-in-numba
@nb.njit('int_(uint32)', inline='always')
def popcount(v):
    v = v - ((v >> 1) & 0x55555555)
    v = (v & 0x33333333) + ((v >> 2) & 0x33333333)
    c = np.uint32((v + (v >> 4) & 0xF0F0F0F) * 0x1010101) >> 24
    return c

# Count the number of bits in a 64-bit integer
@nb.njit(inline='always')
def bit_count(n):
    if n < (1 << 30):
        return popcount(np.uint32(n))
    else:
        return popcount(np.uint32(n)) + popcount(np.uint32(n >> 32))

# Mutate `out` so not to create an expensive new temporary array
@nb.njit
def int_to_path(n, out, a, b):
    for i in range(len(out)):
        out[i] = a[i] if ((n >> i) & 1) else b[i]

@nb.njit(['(int32[:], int32[:], int64, int64)', '(int64[:], int64[:], int64, int64)'])
def traverse_fast(a, b, n_a, n_b):
    # This assertion is needed because the paths are encoded using 64-bit.
    # This should not be a problem in practice since the number of solutions to
    # test would be impracticably huge to test using this algorithm anyway.
    assert n_a + n_b < 62

    max_iter = 1 << (n_a + n_b)
    path = np.empty(n_a + n_b, dtype=a.dtype)
    score, best_score, best_i = 0, 0, 0

    # Iterate over all cases (more than the set of possible solution)
    for i in range(max_iter):
        # Filter the possible solutions
        if bit_count(i) != n_b:
            continue

        # Analyse the score of the solution
        int_to_path(i, path, a, b)
        score = count_unique(path)

        # Store it if it better than the previous one
        if score > best_score:
            best_score = score
            best_i = i

    int_to_path(best_i, path, a, b)
    return best_score, path

此实现在我机器上的大小为 8 的数组上大约快 30 倍。 On 可以使用多个内核来进一步加速。但是，我认为最好专注于寻找次指数的实现，以避免浪费更多的计算资源。请注意，路径与初始函数不同，但随机数组上的分数相同。它可以帮助其他人在更大的阵列上测试他们的实现，而无需等待很长时间。