Repository: element-plus/element-plus
File: packages/components/tree-v2/src/composables/useTree.ts (TypeScript)

Code before:
      expandNode(node)
    }
  }

  function setExpandedKeys(keys: TreeKey[]) {
    const expandedKeys = new Set<TreeKey>()
    const nodeMap = tree.value!.treeNodeMap
    const keysSet = new Set(keys)

    for (const [key] of nodeMap) {
      let node = nodeMap.get(key)
      while (node) {
        if (keysSet.has(node.key)) {
          if (!expandedKeys.has(node.key)) {
            expandedKeys.add(node.key)
            node.expanded = true
            node.parent && keysSet.add(node.parent.key)
          }
        } else {
          node.expanded = false
        }
        node = node.parent
      }
    }

    expandedKeySet.value = expandedKeys
  }

  function handleNodeClick(node: TreeNode, e: MouseEvent) {
    emit(NODE_CLICK, node.data, node, e)
    handleCurrentChange(node)
    if (props.expandOnClickNode) {
      toggleExpand(node)
    }
    if (
      props.showCheckbox &&
      (props.checkOnClickNode || (node.isLeaf && props.checkOnClickLeaf)) &&
      !node.disabled
    ) {
      toggleCheckbox(node, !isChecked(node), true)
    }
  }

  function handleNodeDrop(node: TreeNode, e: DragEvent) {
    emit(NODE_DROP, node.data, node, e)
  }

  function handleCurrentChange(node: TreeNode) {
    if (!isCurrent(node)) {
      currentKey.value = node.key
      emit(CURRENT_CHANGE, node.data, node)

Review comment (suggestion):
The previous logic was successful because `expandedKeySet` was fully reassign and there was no `expanded` property to handle in `node`.

Now we can replicate the logic by resetting the previous nodes expanded before the assign logic.

Pseudo code:
```diff
  function setExpandedKeys(keys: TreeKey[]) {
    const expandedKeys = new Set<TreeKey>()
    const nodeMap = tree.value!.treeNodeMap
+    expandedKeySet.value.forEach((key) => {
+      const node = nodeMap.get(key)!
+      expandedKeySet.value.delete(node.key)
+      node.expanded = false
+    })

    keys.forEach((k) => {
      let node = nodeMap.get(k)
      while (node && !expandedKeys.has(node.key)) {
        expandedKeys.add(node.key)
        node.expanded = true
        node = node.parent
      }
    })

    expandedKeySet.value = expandedKeys
  }
```

Code after:
      expandNode(node)
    }
  }

  function setExpandedKeys(keys: TreeKey[]) {
    const expandedKeys = new Set<TreeKey>()
    const nodeMap = tree.value!.treeNodeMap

    expandedKeySet.value.forEach((key) => {
      const node = nodeMap.get(key)!
      expandedKeySet.value.delete(node.key)
      node.expanded = false
    })

    keys.forEach((k) => {
      let node = nodeMap.get(k)
      while (node && !expandedKeys.has(node.key)) {
        expandedKeys.add(node.key)
        node.expanded = true
        node = node.parent
      }
    })

    expandedKeySet.value = expandedKeys
  }

  function handleNodeClick(node: TreeNode, e: MouseEvent) {
    emit(NODE_CLICK, node.data, node, e)
    handleCurrentChange(node)
    if (props.expandOnClickNode) {
      toggleExpand(node)
    }
    if (
      props.showCheckbox &&
      (props.checkOnClickNode || (node.isLeaf && props.checkOnClickLeaf)) &&
      !node.disabled
    ) {
      toggleCheckbox(node, !isChecked(node), true)
    }
  }

  function handleNodeDrop(node: TreeNode, e: DragEvent) {
    emit(NODE_DROP, node.data, node, e)
  }

  function handleCurrentChange(node: TreeNode) {
    if (!isCurrent(node)) {
      currentKey.value = node.key
      emit(CURRENT_CHANGE, node.data, node)
    }
  }