【Python】NumPyのstackとsqueezeの違いを解説

はじめに

データ分析や機械学習でNumPyやPyTorchを使っていると、よく出てくるのがstackとsqueeze。
今まで，AIに修正させていたので，具体的にどう違うのかわからず使っていました．
そこで，この記事では、両者の違いを実際のコード例を交えて整理しようと思います。

stackとは？

👉 新しい次元を追加して配列を結合する関数 です。

import numpy as np

a = np.array([1, 2, 3])  # shape: (3,)
b = np.array([4, 5, 6])  # shape: (3,)

stacked = np.stack([a, b], axis=0)
print(stacked.shape)  # (2, 3)
print(stacked)
# [[1 2 3]
#  [4 5 6]]

stacked2 = np.stack([a, b], axis=1)
print(stacked2.shape)  # (3, 2)
print(stacked2)
# [[1 4]
#  [2 5]
#  [3 6]]

axis=0 → 行方向に積み上げる
axis=1 → 列方向に積み上げる

squeezeとは？

👉 サイズ1の次元を削除する関数 です。

x = np.array([[[1], [2], [3]]])
print(x.shape)  # (1, 3, 1)

y = np.squeeze(x)
print(y.shape)  # (3,)
print(y)        # [1 2 3]

(1, 3, 1) → (3,) のように、余計な「1次元」を消してくれる

実際のケース：DataFrameからの変換

例えば次のようなデータがあるとします。

import pandas as pd
import numpy as np

df = pd.DataFrame({
    "before_predicted_position": [
        np.array([1, 2, 3]),
        np.array([4, 5, 6]),
        np.array([7, 8, 9]),
        np.array([10, 11, 12]),
        np.array([13, 14, 15]),
    ]
})

このとき、

print(df["before_predicted_position"].to_numpy().shape)   # (5,)
print(df["before_predicted_position"].to_numpy()[0].shape)  # (3,)

一見 (5, 3) に見えそうですが、実際は「長さ5のobject配列」になっています。

そこで stack の出番です👇

arr = np.stack(df["before_predicted_position"].to_numpy())
print(arr.shape)  # (5, 3)
print(arr)
# [[ 1  2  3]
#  [ 4  5  6]
#  [ 7  8  9]
#  [10 11 12]
#  [13 14 15]]

これで 綺麗に (5,3) の数値配列 に整形できました ✅

まとめ

stack → 新しい次元を作って配列を結合
squeeze → 不要な次元（サイズ1）を削除

💡 個人的には、
「配列を揃えてまとめたい」→ stack
「配列をスリムにしたい」→ squeeze
とすればわかりやすいかもです。