数组的最大与和

标签: 位运算 数组 动态规划 状态压缩

难度: Hard

给你一个长度为 n 的整数数组 nums 和一个整数 numSlots ,满足2 * numSlots >= n 。总共有 numSlots 个篮子,编号为 1 到 numSlots 。

你需要把所有 n 个整数分到这些篮子中,且每个篮子 至多 有 2 个整数。一种分配方案的 与和 定义为每个数与它所在篮子编号的 按位与运算 结果之和。

  • 比方说,将数字 [1, 3] 放入篮子 1 中,[4, 6] 放入篮子 2 中,这个方案的与和为 (1 AND 1) + (3 AND 1) + (4 AND 2) + (6 AND 2) = 1 + 1 + 0 + 2 = 4 。

请你返回将 nums 中所有数放入 numSlots 个篮子中的最大与和。

示例 1:

输入:nums = [1,2,3,4,5,6], numSlots = 3
输出:9
解释:一个可行的方案是 [1, 4] 放入篮子 1 中,[2, 6] 放入篮子 2 中,[3, 5] 放入篮子 3 中。
最大与和为 (1 AND 1) + (4 AND 1) + (2 AND 2) + (6 AND 2) + (3 AND 3) + (5 AND 3) = 1 + 0 + 2 + 2 + 3 + 1 = 9 。

示例 2:

输入:nums = [1,3,10,4,7,1], numSlots = 9
输出:24
解释:一个可行的方案是 [1, 1] 放入篮子 1 中,[3] 放入篮子 3 中,[4] 放入篮子 4 中,[7] 放入篮子 7 中,[10] 放入篮子 9 中。
最大与和为 (1 AND 1) + (1 AND 1) + (3 AND 3) + (4 AND 4) + (7 AND 7) + (10 AND 9) = 1 + 1 + 3 + 4 + 7 + 8 = 24 。
注意,篮子 2 ,5 ,6 和 8 是空的,这是允许的。

提示:

  • n == nums.length
  • 1 <= numSlots <= 9
  • 1 <= n <= 2 * numSlots
  • 1 <= nums[i] <= 15

Submission

运行时间: 35 ms

内存: 16.3 MB

class Edge:
    def __init__(self, to, rev, cap, cost):
        self.to = to
        self.rev = rev
        self.cap = cap
        self.cost = cost

def add_edge(graph, from_, to, cap, cost):
    graph[from_].append(Edge(to, len(graph[to]), cap, cost))
    graph[to].append(Edge(from_, len(graph[from_])-1, 0, -cost))

def spfa(graph, start, end, dist, parent):
    inf = 10**9
    for i in range(len(dist)):
        dist[i] = inf
    dist[start] = 0
    in_queue = [False] * len(graph)
    in_queue[start] = True
    queue = deque([start])
    while queue:
        v = queue.popleft()
        in_queue[v] = False
        for i, edge in enumerate(graph[v]):
            if edge.cap > 0 and dist[v] + edge.cost < dist[edge.to]:
                dist[edge.to] = dist[v] + edge.cost
                parent[edge.to] = (v, i)
                if not in_queue[edge.to]:
                    queue.append(edge.to)
                    in_queue[edge.to] = True
    return dist[end] != inf

from collections import deque

class Solution:

    def maximumANDSum(self, nums, num_slots):
        n, m = len(nums), num_slots
        graph = [[] for _ in range(n + m + 2)]
        start, end = n + m, n + m + 1
        inf = 10**9

        for i, num in enumerate(nums):
            add_edge(graph, start, i, 1, 0)
            for j in range(1, m + 1):
                add_edge(graph, i, n + j - 1, inf, -(num & j))

        for i in range(m):
            add_edge(graph, n + i, end, 2, 0)

        ans = 0
        dist = [0] * len(graph)
        parent = [None] * len(graph)
        while spfa(graph, start, end, dist, parent):
            min_flow = inf
            v = end
            while v != start:
                p, i = parent[v]
                min_flow = min(min_flow, graph[p][i].cap)
                v = p
            v = end
            while v != start:
                p, i = parent[v]
                graph[p][i].cap -= min_flow
                graph[v][graph[p][i].rev].cap += min_flow
                v = p
            ans -= dist[end] * min_flow
        return ans

Explain

该题解使用了网络流的最小费用最大流算法来解决问题。首先,构建一个图,其中包含一个源点和一个汇点,以及代表数字和篮子的节点。为每个数字到每个篮子的边设置一个容量为无穷大,费用为该数字与篮子编号的按位与的负值。这样设置费用是为了最小费用流算法能求得最大与和。同时,将所有数字连接到源点的边的容量设为1,费用设为0,意味着每个数字只能选择一个篮子。每个篮子到汇点的边的容量设置为2,费用为0,表示每个篮子最多只能放两个数字。通过不断找到一条从源点到汇点的最小费用的路径并增加流量,直到没有增广路径存在,计算总的最小费用,并将其转化为最大与和。

时间复杂度: O(V^2*E)

空间复杂度: O(V + E)

from collections import deque

class Edge:
    def __init__(self, to, rev, cap, cost):
        # 初始化边,包含终点、反向边索引、容量和费用
        self.to = to
        self.rev = rev
        self.cap = cap
        self.cost = cost

def add_edge(graph, from_, to, cap, cost):
    # 向图中添加边
    graph[from_].append(Edge(to, len(graph[to]), cap, cost))
    graph[to].append(Edge(from_, len(graph[from_]-1), 0, -cost))

def spfa(graph, start, end, dist, parent):
    # SPFA算法寻找最小费用路径
    inf = 10**9
    for i in range(len(dist)):
        dist[i] = inf
    dist[start] = 0
    in_queue = [False] * len(graph)
    in_queue[start] = True
    queue = deque([start])
    while queue:
        v = queue.popleft()
        in_queue[v] = False
        for i, edge in enumerate(graph[v]):
            if edge.cap > 0 and dist[v] + edge.cost < dist[edge.to]:
                dist[edge.to] = dist[v] + edge.cost
                parent[edge.to] = (v, i)
                if not in_queue[edge.to]:
                    queue.append(edge.to)
                    in_queue[edge.to] = True
    return dist[end] != inf

class Solution:

    def maximumANDSum(self, nums, num_slots):
        # 初始化图和变量
        n, m = len(nums), num_slots
        graph = [[] for _ in range(n + m + 2)]
        start, end = n + m, n + m + 1
        inf = 10**9
        # 添加边
        for i, num in enumerate(nums):
            add_edge(graph, start, i, 1, 0)
            for j in range(1, m + 1):
                add_edge(graph, i, n + j - 1, inf, -(num & j))
        for i in range(m):
            add_edge(graph, n + i, end, 2, 0)
        # 最小费用最大流主循环
        ans = 0
        dist = [0] * len(graph)
        parent = [None] * len(graph)
        while spfa(graph, start, end, dist, parent):
            min_flow = inf
            v = end
            while v != start:
                p, i = parent[v]
                min_flow = min(min_flow, graph[p][i].cap)
                v = p
            v = end
            while v != start:
                p, i = parent[v]
                graph[p][i].cap -= min_flow
                graph[v][graph[p][i].rev].cap += min_flow
                v = p
            ans -= dist[end] * min_flow
        return ans

Explore

在该问题中,目标是最大化所有数字与其对应篮子编号的按位与的和。为了使用最小费用最大流算法解决这一最大化问题,我们将每个按位与的结果转化为负值作为费用。这样,原本求最小费用的算法就可以通过寻找总费用最小的方式来实际上找到按位与和的最大值。因此,通过设置费用为按位与结果的负值,我们将问题转化为标准的最小费用流问题,使得算法可以应用于此场景。

SPFA(Shortest Path Faster Algorithm)算法是Bellman-Ford算法的一种改进,它在实际应用中通常比Bellman-Ford更高效。与Dijkstra算法相比,SPFA算法的主要优势在于它可以正确处理图中存在负权边的情况。在最小费用流问题中,由于按位与操作结果转化为负值费用,负权边是常见的。因此,SPFA算法更适合本题解中的场景,能够有效地寻找最小费用路径,即使路径中包含负权边。

在构建网络流模型时,每个篮子到汇点的边的容量设置为2意味着通过这条边的流量不能超过2。这直接限制了每个篮子节点可以接受的总流量,即最多只能有两个数字的流量(每个数字对应单位流量)流入汇点。因此,这种容量设置确保了每个篮子在解决方案中最多只能包含两个数字。

在算法实现中,虽然使用的是最小费用流算法,但是每个数字到篮子的边的费用被设置为了该数字与篮子编号按位与结果的负值。因此,最小化这些负值费用的和实际上等同于最大化原始的按位与结果的和。在算法的最后,通过计算所有路径的负费用乘以流量的总和,并取其负值(即`-dist[end] * min_flow`),从而得到最大的与和。这样确保了最终输出是最大与和而非最小值。