【编程题】最小农田浇水费用

题目描述

一块农田又n块地，为了能给n块地顺利的浇上水，对于每块地可以选择在本地打井，也可以修管道从另一块地（这块地可能打了井，它的水也是从其他地引过来的）将水引过来。
当然打井和修管道的费用有差别。已知在第i块地打井需要耗费wi元，在第i、j间修管道要pij元。由于地较多，需要选择在哪些地打井，哪些地之间修管道。请编程做决策，求出最小费用。

输入输出描述

以下n行（第一个输入），一次给出整数w1, …, wn（每块地的打井费用）。紧接着是一个n*n的整数矩阵，矩阵的第i行第j列的数代表pij（两块地之间建立管道的费用）。每行的两个数之间有一个空格隔开。

6
5
4
4
3
1
20
0 2 2 2 9 9
2 0 3 3 9 9
2 3 0 4 9 9
2 3 4 0 9 9
9 9 9 9 0 9
9 9 9 9 9 0

输出最小费用值

16

代码

初步思路及实现，有错误的地方烦请指正！

def func(dig_prices, matrix):
    min_price = min(dig_prices)
    min_connected = float('inf')
    idx = [0, 0]
    for i in range(n):
        for j in range(n):
            if matrix[i][j] < min_price and matrix[i][j] < min_connected:
                min_connected = matrix[i][j]
                idx[0], idx[1] = i, j

    visited = []
    res = 0
    if min_connected != float('inf'):
        visited.append(idx[0])
        visited.append(idx[1])
        res += matrix[idx[0]][idx[1]]
        for x in range(n):
            if x in visited:
                continue
            else:
                index, dig = compare_price(x, visited, dig_prices, matrix)
                if dig:
                    res += dig_prices[x]
                    visited.append(x)
                else:
                    res += matrix[x][index]
                    visited.append(index)
    else:
        visited.append(dig_prices.index(min_price))
        res += min_price
        for x in range(n):
            if x in visited:
                continue
            else:
                index, dig = compare_price(x, visited, dig_prices, matrix)
                if dig:
                    res += dig_prices[x]
                    visited.append(x)
                else:
                    res += matrix[x][index]
                    visited.append(index)
    return res


def compare_price(i, visit, dig_price, mat):
    dig_cost = dig_price[i]
    flag = True
    idx = -1
    min_cost = float('inf')
    for j in visit:
        if mat[i][j] < dig_cost and mat[i][j] < min_cost:
            idx = j
            flag = False
            min_cost = mat[i][j]

    return idx, flag


n = int(input())
dig_prices = []
for _ in range(n):
    dig_prices.append(int(input()))

matrix = []
for _ in range(n):
    matrix.append(list(map(int, input().split())))

# print(dig_prices)
# print(matrix)

res = func(dig_prices, matrix)
print(res)