Question

1167. Minimum Cost to Connect Sticks

中文文档

Description

You have some number of sticks with positive integer lengths. These lengths are given as an array sticks, where sticks[i] is the length of the ith stick.

You can connect any two sticks of lengths x and y into one stick by paying a cost of x + y. You must connect all the sticks until there is only one stick remaining.

Return _the minimum cost of connecting all the given sticks into one stick in this way_.

 

Example 1:

Input: sticks = [2,4,3]
Output: 14
Explanation: You start with sticks = [2,4,3].
1. Combine sticks 2 and 3 for a cost of 2 + 3 = 5. Now you have sticks = [5,4].
2. Combine sticks 5 and 4 for a cost of 5 + 4 = 9. Now you have sticks = [9].
There is only one stick left, so you are done. The total cost is 5 + 9 = 14.

Example 2:

Input: sticks = [1,8,3,5]
Output: 30
Explanation: You start with sticks = [1,8,3,5].
1. Combine sticks 1 and 3 for a cost of 1 + 3 = 4. Now you have sticks = [4,8,5].
2. Combine sticks 4 and 5 for a cost of 4 + 5 = 9. Now you have sticks = [9,8].
3. Combine sticks 9 and 8 for a cost of 9 + 8 = 17. Now you have sticks = [17].
There is only one stick left, so you are done. The total cost is 4 + 9 + 17 = 30.

Example 3:

Input: sticks = [5]
Output: 0
Explanation: There is only one stick, so you don't need to do anything. The total cost is 0.

 

Constraints:

• 1 <= sticks.length <= 104
• 1 <= sticks[i] <= 104

Solutions

Solution 1: Greedy + Priority Queue (Min Heap)

We can use a greedy approach, each time choosing the shortest two sticks to connect, which ensures the minimum cost of connection.

Therefore, we can use a priority queue (min heap) to maintain the current stick lengths. Each time, we take out two sticks from the priority queue to connect, then put the connected stick back into the priority queue, until there is only one stick left in the priority queue.

The time complexity is $O(n \times \log n)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array sticks.

class Solution:
  def connectSticks(self, sticks: List[int]) -> int:
    heapify(sticks)
    ans = 0
    while len(sticks) > 1:
      z = heappop(sticks) + heappop(sticks)
      ans += z
      heappush(sticks, z)
    return ans

class Solution {
  public int connectSticks(int[] sticks) {
    PriorityQueue<Integer> pq = new PriorityQueue<>();
    for (int x : sticks) {
      pq.offer(x);
    }
    int ans = 0;
    while (pq.size() > 1) {
      int z = pq.poll() + pq.poll();
      ans += z;
      pq.offer(z);
    }
    return ans;
  }
}

class Solution {
public:
  int connectSticks(vector<int>& sticks) {
    priority_queue<int, vector<int>, greater<int>> pq;
    for (auto& x : sticks) {
      pq.push(x);
    }
    int ans = 0;
    while (pq.size() > 1) {
      int x = pq.top();
      pq.pop();
      int y = pq.top();
      pq.pop();
      int z = x + y;
      ans += z;
      pq.push(z);
    }
    return ans;
  }
};

func connectSticks(sticks []int) (ans int) {
  hp := &hp{sticks}
  heap.Init(hp)
  for hp.Len() > 1 {
    x, y := heap.Pop(hp).(int), heap.Pop(hp).(int)
    ans += x + y
    heap.Push(hp, x+y)
  }
  return
}

type hp struct{ sort.IntSlice }

func (h hp) Less(i, j int) bool { return h.IntSlice[i] < h.IntSlice[j] }
func (h *hp) Push(v any)    { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() any {
  a := h.IntSlice
  v := a[len(a)-1]
  h.IntSlice = a[:len(a)-1]
  return v
}

function connectSticks(sticks: number[]): number {
  const pq = new Heap(sticks);
  let ans = 0;
  while (pq.size() > 1) {
    const x = pq.pop();
    const y = pq.pop();
    ans += x + y;
    pq.push(x + y);
  }
  return ans;
}

type Compare<T> = (lhs: T, rhs: T) => number;

class Heap<T = number> {
  data: Array<T | null>;
  lt: (i: number, j: number) => boolean;
  constructor();
  constructor(data: T[]);
  constructor(compare: Compare<T>);
  constructor(data: T[], compare: Compare<T>);
  constructor(data: T[] | Compare<T>, compare?: (lhs: T, rhs: T) => number);
  constructor(
    data: T[] | Compare<T> = [],
    compare: Compare<T> = (lhs: T, rhs: T) => (lhs < rhs ? -1 : lhs > rhs ? 1 : 0),
  ) {
    if (typeof data === 'function') {
      compare = data;
      data = [];
    }
    this.data = [null, ...data];
    this.lt = (i, j) => compare(this.data[i]!, this.data[j]!) < 0;
    for (let i = this.size(); i > 0; i--) this.heapify(i);
  }

  size(): number {
    return this.data.length - 1;
  }

  push(v: T): void {
    this.data.push(v);
    let i = this.size();
    while (i >> 1 !== 0 && this.lt(i, i >> 1)) this.swap(i, (i >>= 1));
  }

  pop(): T {
    this.swap(1, this.size());
    const top = this.data.pop();
    this.heapify(1);
    return top!;
  }

  top(): T {
    return this.data[1]!;
  }
  heapify(i: number): void {
    while (true) {
      let min = i;
      const [l, r, n] = [i * 2, i * 2 + 1, this.data.length];
      if (l < n && this.lt(l, min)) min = l;
      if (r < n && this.lt(r, min)) min = r;
      if (min !== i) {
        this.swap(i, min);
        i = min;
      } else break;
    }
  }

  clear(): void {
    this.data = [null];
  }

  private swap(i: number, j: number): void {
    const d = this.data;
    [d[i], d[j]] = [d[j], d[i]];
  }
}