python - 求 c 使得 sum(x+c) over positives = K

Question

假设我有一个在 Python 中具有正值和负值的一维数组 x，例如:

x = random.rand(10) * 10

对于 K 的给定正值，我想找到使总和的偏移量c > 数组 y = x + c 的 positive 元素等于 K。

我怎样才能有效地解决这个问题？

Answer 1

如何通过二分搜索来确定 x + c 的哪些元素将对总和做出贡献，然后求解线性方程？这段代码的运行时间是 O(n log n)，但在 Python 中只完成了 O(log n) 的工作。通过更复杂的分区策略可以将运行时间降低到 O(n)。我不确定是否会带来实际的改进。

import numpy as np

def findthreshold(x, K):
    x = np.sort(np.array(x))[::-1]
    z = np.cumsum(np.array(x))
    l = 0
    u = x.size
    while u - l > 1:
        m = (l + u) // 2
        if z[m] - (m + 1) * x[m] >= K:
            u = m
        else:
            l = m
    return (K - z[l]) / (l + 1)

def test():
    x = np.random.rand(10)
    K = np.random.rand() * x.size
    c = findthreshold(x, K)
    assert np.abs(K - np.sum(np.clip(x + c, 0, np.inf))) / K <= 1e-8

这是一个随机化的预期 O(n) 变体。它更快(在我的机器上，对于大输入)，但不是那么快。当心两个版本中的灾难性取消。

def findthreshold2(x, K):
    sumincluded = 0
    includedsize = 0
    while x.size > 0:
        pivot = x[np.random.randint(x.size)]
        above = x[x > pivot]
        if sumincluded + np.sum(above) - (includedsize + above.size) * pivot >= K:
            x = above
        else:
            notbelow = x[x >= pivot]
            sumincluded += np.sum(notbelow)
            includedsize += notbelow.size
            x = x[x < pivot]
    return (K - sumincluded) / includedsize