f163. 貨物分配

解題:

我們先介紹用到的資料結構:

• tempw：暫時存放底層貨櫃的重量
• inbox：要放入的貨物重量
• parent：整個貨運站的樹的上一個父節點(若為-1則為根，並且node從1開始算)
• child：整個貨運站的樹的節點的子節點(若長度為空為leaf)
• leaf：存放葉子的節點方便查詢並且建立w_t
• w：整個貨運站的每個節點的重量

dfs函數講解

我們利用dfs函式去對於每個leaf($n<=i<2*n$)去做查詢，先把$wt_{nownode}$增加$tempw_{i-n}$，並且把nownode設為舊的nownode的父節點(此行為到root為止)

最後回傳wt作為重量的原始值

solve函式講解

我們利用solve函式對inbox內所有要進行分配的貨物進行遍歷，並且把nownode的初始值設為root(方便進行dfs)，接著進入while loop，如果相等或左邊節點重量小於右邊節點重量就將貨物流入左邊節點並且把此節點的重量+上流入貨物的重量(wi)，則相反的如果右邊的節點重量小於左邊的節點重量就流入右邊的節點並且把此流入節點的重量+上流入貨物的重量(wi)，這個迴圈到leaf($child_{nownode}$的長度為0)把這個leaf加進去out陣列，然後跳出迴圈查看下一個貨物

AC Code

python

from sys import stdin
def dfs(leaf):
    wt=[0]*(2*n)
    root=0
    for i in leaf:
        wi=tempw[i-n]
        now_node=i
        while nownode!=-1:
            wt[nownode]+=wi
            nownode=parent[nownode]
        root=nownode
    return wt
def solve(w,r):
    out=[]
    for i in inbox:
        nownode=r
        w[nownode]+=i
        while True:
            if len(child[nownode])==0:
                out.append(nownode)
                break
            if w[child[nownode][0]]<=w[child[nownode][1]]:
                w[child[nownode][0]]+=i
                now_node=child[nownode][0]
            elif w[child[nownode][0]]>w[child[nownode][1]]:
                w[child[nownode][1]]+=i
                nownode=child[nownode][1]
    return out
n,m=map(int,stdin.readline().strip().split())
tempw=list(map(int,stdin.readline().strip().split()))
inbox=list(map(int,stdin.readline().strip().split()))
parent=[-1]*(2*n)
child=[[] for _ in range(2*n)]
for _ in range(n-1):
    a,b,c=map(int,stdin.readline().strip().split())
    parent[b]=a
    parent[c]=a
    child[a].append(b)
    child[a].append(c)
leaf=[i for i in range(n,2*n)]
r=0
for i in range(1,2*n):
    if parent[i]==-1:
        r=i
        break
print(solve(dfs(leaf),r))

C++

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
int n,m;
ll w[200005];
int parent[200005];
vector<int> child[200005];
int temp_w[100005];
int inbox[100005]; 

void init_w(){
    for(int i=0;i<2*n;i++){
        ll weight=temp_w[i-n];
        int now_node=i;
        while(now_node!=-1){
            w[now_node]+=weight;
            now_node=parent[now_node];
        }
    }
}

void solve(int root){
    for(int i=0;i<m;++i){
        int weight=inbox[i];
        int now_node=root;
        w[now_node]+=weight;
        while(true){
            if(child[now_node].empty()){
                cout<<now_node<<(i==m-1 ? "" : " ");
                break;
            }
            int left=child[now_node][0];
            int right=child[now_node][1];
            if(w[left]<=w[right]){
                w[left]+=weight;
                now_node=left;
            }else{
                w[right]+=weight;
                now_node=right;
            }
        }
    }
    cout<<endl;
}
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);

    cin>>n>>m;
    for(int i=0;i<2*n;i++) parent[i]=-1;
    for(int i=0;i<n;i++) cin>>temp_w[i];
    for(int i=0;i<m;i++) cin>>inbox[i];
    for(int i=0;i<n-1;i++){
        int u,v,k;
        cin>>u>>v>>k;
        parent[v]=u;
        parent[k]=u;
        child[u].push_back(v);
        child[u].push_back(k);
    }
    int root=1;
    for(int i=1;i<2*n;i++){
        if(parent[i]==-1 && !child[i].empty()){
            root=i;
            break;
        }
    }
    init_w();
    solve(root);
}